# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Checkpoint interval against WAL volume in PostgreSQL.

A pgbench-style TPC-B workload (8 clients updating accounts, tellers and
branches) for a fixed time with checkpoint_timeout = 30 s, then 15 min, then
30 s with full_page_writes off: WAL bytes generated, and WAL bytes per
transaction. The load is driven from Python so it needs only the database.
Needs the PostgreSQL service from the compose file (port 55500).
"""
import argparse, json, random, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55500/postgres", help="Connection string.")
p.add_argument("--scale", type=int, default=50, help="pgbench scale (100,000 accounts per unit).")
p.add_argument("--clients", type=int, default=8, help="Client threads.")
p.add_argument("--seconds", type=int, default=120, help="Seconds per setting.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
admin = psycopg.connect(args.dsn, autocommit=True)
A = args.scale * 100_000
admin.execute("DROP TABLE IF EXISTS accounts, tellers, branches, history")
admin.execute("CREATE TABLE branches (bid int PRIMARY KEY, bbalance int, filler char(88))")
admin.execute("CREATE TABLE tellers (tid int PRIMARY KEY, bid int, tbalance int, filler char(84))")
admin.execute("CREATE TABLE accounts (aid int PRIMARY KEY, bid int, abalance int, filler char(84)) WITH (fillfactor = 100)")
admin.execute("CREATE TABLE history (tid int, bid int, aid int, delta int, mtime timestamp, filler char(22))")
admin.execute(f"INSERT INTO branches SELECT b, 0, '' FROM generate_series(1, {args.scale}) b")
admin.execute(f"INSERT INTO tellers SELECT t, (t - 1) / 10 + 1, 0, '' FROM generate_series(1, {args.scale * 10}) t")
admin.execute(f"INSERT INTO accounts SELECT a, (a - 1) / 100000 + 1, 0, '' FROM generate_series(1, {A}) a")
admin.execute("VACUUM ANALYZE")

def workload():
    stop, count = threading.Event(), [0]
    def client():
        con = psycopg.connect(args.dsn, autocommit=False); rnd = random.Random()
        while not stop.is_set():
            aid, tid, bid, d = rnd.randint(1, A), rnd.randint(1, args.scale * 10), rnd.randint(1, args.scale), rnd.randint(-5000, 5000)
            con.execute("UPDATE accounts SET abalance = abalance + %s WHERE aid = %s", (d, aid))
            con.execute("UPDATE tellers SET tbalance = tbalance + %s WHERE tid = %s", (d, tid))
            con.execute("UPDATE branches SET bbalance = bbalance + %s WHERE bid = %s", (d, bid))
            con.execute("INSERT INTO history VALUES (%s, %s, %s, %s, now(), '')", (tid, bid, aid, d))
            con.commit(); count[0] += 1
        con.close()
    ts = [threading.Thread(target=client) for _ in range(args.clients)]
    for t in ts: t.start()
    time.sleep(args.seconds); stop.set()
    for t in ts: t.join()
    return count[0]

rows = []
for label, timeout, fpw in (("checkpoint every 30 s", "30s", "on"), ("checkpoint every 15 min", "15min", "on"),
                            ("every 30 s, full_page_writes off", "30s", "off")):
    admin.execute(f"ALTER SYSTEM SET checkpoint_timeout = '{timeout}'"); admin.execute(f"ALTER SYSTEM SET full_page_writes = {fpw}")
    admin.execute("SELECT pg_reload_conf()"); admin.execute("CHECKPOINT")
    lsn0 = admin.execute("SELECT pg_current_wal_lsn()").fetchone()[0]
    fpi0 = admin.execute("SELECT wal_fpi FROM pg_stat_wal").fetchone()[0]
    tx = workload()
    wal = float(admin.execute("SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), %s)", (lsn0,)).fetchone()[0])
    fpi = admin.execute("SELECT wal_fpi FROM pg_stat_wal").fetchone()[0] - fpi0
    rows.append({"setting": label, "transactions": tx, "wal_mb": wal / 1e6, "wal_bytes_per_tx": wal / tx, "full_page_images": fpi})
    print(f"{label:<34} {tx:>8,} tx   WAL {wal / 1e6:8.1f} MB   {wal / tx:7.0f} bytes per tx   full-page images {fpi:,}")
admin.execute("ALTER SYSTEM RESET checkpoint_timeout"); admin.execute("ALTER SYSTEM RESET full_page_writes"); admin.execute("SELECT pg_reload_conf()")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"settings": rows}}, f)