# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""One old snapshot pins the version chain.

A PostgreSQL table of 1 million accounts updated at random by 8 clients for
two minutes with autovacuum on; the table size and n_dead_tup afterwards. The
same again while another session holds a Repeatable Read snapshot open. Then
that session commits, VACUUM runs, and n_dead_tup and the size are read again.
Needs the PostgreSQL service from the compose file (port 55802).
"""
import argparse, json, random, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55802/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=1_000_000, help="Accounts.")
p.add_argument("--seconds", type=int, default=120, help="Update time per run.")
p.add_argument("--clients", type=int, default=8, help="Updating clients.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

admin = psycopg.connect(args.dsn, autocommit=True)
admin.execute("ALTER SYSTEM SET autovacuum_naptime = '5s'"); admin.execute("SELECT pg_reload_conf()")
def reset():
    admin.execute("DROP TABLE IF EXISTS accounts")
    admin.execute(f"CREATE TABLE accounts AS SELECT g AS id, 0 AS balance, repeat('x', 80) AS filler FROM generate_series(1, {args.rows}) g")
    admin.execute("ALTER TABLE accounts ADD PRIMARY KEY (id)"); admin.execute("VACUUM ANALYZE accounts")
def load():
    stop = time.time() + args.seconds; n = [0]
    def worker(seed):
        rng = random.Random(seed); c = psycopg.connect(args.dsn, autocommit=True)
        while time.time() < stop:
            c.execute("UPDATE accounts SET balance = balance + 1 WHERE id = %s", (rng.randrange(1, args.rows + 1),)); n[0] += 1
    ts = [threading.Thread(target=worker, args=(i,)) for i in range(args.clients)]
    for t in ts: t.start()
    for t in ts: t.join()
    return n[0]
def stat():
    admin.execute("SELECT pg_stat_force_next_flush()"); time.sleep(1)
    size = admin.execute("SELECT pg_total_relation_size('accounts')").fetchone()[0] / 2**20
    dead = admin.execute("SELECT n_dead_tup FROM pg_stat_user_tables WHERE relname = 'accounts'").fetchone()[0]
    return size, dead
rows = []
reset(); s0, _ = stat(); n = load(); s, d = stat()
rows.append({"run": "no open snapshot", "updates": n, "size_mb_before": s0, "size_mb_after": s, "n_dead_tup": d})
reset(); s0, _ = stat()
old = psycopg.connect(args.dsn); old.execute("BEGIN ISOLATION LEVEL REPEATABLE READ"); old.execute("SELECT count(*) FROM accounts").fetchone()
n = load(); s, d = stat()
rows.append({"run": "one Repeatable Read snapshot left open", "updates": n, "size_mb_before": s0, "size_mb_after": s, "n_dead_tup": d})
old.execute("COMMIT"); old.close(); admin.execute("VACUUM accounts"); s, d = stat()
rows.append({"run": "after the old session commits, then VACUUM", "updates": 0, "size_mb_before": None, "size_mb_after": s, "n_dead_tup": d})
for r in rows: print(f"{r['run']:<44} updates {r['updates']:>8,}  size {r['size_mb_after']:7.1f} MB  n_dead_tup {r['n_dead_tup']:>9,}")

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