# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The savepoint cliff in PostgreSQL: more than 64 subtransactions.

Session A opens a transaction that creates N subtransactions (an EXCEPTION
block per UPDATE in a loop) and leaves it open. Meanwhile pgbench runs a
read-only query on the same rows with 8 clients for 20 seconds. N = 0, 50,
64, 65 and 200: read throughput against N.
Needs the PostgreSQL service from the compose file (port 55806) and the
docker CLI (pgbench runs inside the container).
"""
import argparse, json, re, subprocess
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55806/postgres", help="Connection string.")
p.add_argument("--container", default="savepoint-cliff-in-postgres-more-than-postgres-1", help="PostgreSQL container (for pgbench).")
p.add_argument("--subxacts", default="0,50,64,65,200", help="Subtransactions in the open transaction.")
p.add_argument("--seconds", type=int, default=20, help="pgbench duration.")
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("DROP TABLE IF EXISTS wallets"); admin.execute("CREATE TABLE wallets (owner text PRIMARY KEY, balance int)")
admin.execute("INSERT INTO wallets SELECT 'w' || g, 0 FROM generate_series(1, 1000) g"); admin.execute("INSERT INTO wallets VALUES ('anna', 0)")
subprocess.run(["docker", "exec", args.container, "sh", "-c", "echo \"SELECT sum(balance) FROM wallets WHERE owner IN ('anna', 'w1', 'w2');\" > /tmp/select_wallets.sql"], check=True)
rows = []
for n in map(int, args.subxacts.split(",")):
    a = psycopg.connect(args.dsn); a.execute("BEGIN")
    a.execute("UPDATE wallets SET balance = balance + 1 WHERE owner = 'w1'")                       # the top-level transaction writes too
    if n:
        a.execute(f"""DO $$ BEGIN FOR i IN 1..{n} LOOP
                        BEGIN UPDATE wallets SET balance = balance + 1 WHERE owner = 'anna';
                        EXCEPTION WHEN OTHERS THEN NULL; END; END LOOP; END $$""")              # each EXCEPTION block is a subtransaction
    out = subprocess.run(["docker", "exec", args.container, "pgbench", "-U", "postgres", "-n", "-T", str(args.seconds), "-c", "8", "-j", "4",
                          "-f", "/tmp/select_wallets.sql", "postgres"], capture_output=True, text=True)
    tps = float(re.search(r"tps = ([\d.]+)", out.stdout).group(1))
    rows.append({"subtransactions": n, "read_tps": tps}); a.rollback(); a.close()
base = rows[0]["read_tps"]
for r in rows: r["vs_none"] = r["read_tps"] / base
for r in rows: print(f"{r['subtransactions']:>4} subtransactions open: {r['read_tps']:9.0f} reads/s ({r['vs_none']:.2f}x of none)")

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