# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Write skew under Repeatable Read, and the two cures.

Two transactions each check that an owner's total covers a withdrawal, then
take 80 from different accounts. Under Repeatable Read (snapshot isolation)
both commit and the total goes negative. Then Serializable, and Repeatable
Read with the conflict materialized on a shared row.
Needs the PostgreSQL service from the compose file (port 55461).
"""
import argparse, json, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55461/postgres", help="Connection string.")
p.add_argument("--withdraw", type=int, default=80, help="Amount each transaction takes.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

admin = psycopg.connect(args.dsn, autocommit=True)
def reset():
    admin.execute("DROP TABLE IF EXISTS acct, owners")
    admin.execute("CREATE TABLE acct(owner int, kind text, balance int, PRIMARY KEY(owner, kind))")
    admin.execute("INSERT INTO acct VALUES (1,'main',60), (1,'savings',40)")
    admin.execute("CREATE TABLE owners(id int PRIMARY KEY, version int)")
    admin.execute("INSERT INTO owners VALUES (1, 0)")
total = lambda: admin.execute("SELECT sum(balance) FROM acct WHERE owner = 1").fetchone()[0]


def run(level, materialize):
    reset()
    a, b = psycopg.connect(args.dsn), psycopg.connect(args.dsn)
    out = {"A": "committed", "B": "committed"}
    for c in (a, b): c.execute(f"BEGIN ISOLATION LEVEL {level}")
    if materialize:
        a.execute("UPDATE owners SET version = version + 1 WHERE id = 1")
    seen = [c.execute("SELECT sum(balance) FROM acct WHERE owner = 1").fetchone()[0] for c in (a, b)]
    a.execute("UPDATE acct SET balance = balance - %s WHERE owner = 1 AND kind = 'main'", (args.withdraw,))
    def b_side():
        try:
            if materialize:                              # B touches the same row: it waits for A, then fails
                b.execute("UPDATE owners SET version = version + 1 WHERE id = 1")
            b.execute("UPDATE acct SET balance = balance - %s WHERE owner = 1 AND kind = 'savings'", (args.withdraw,))
            b.commit()
        except psycopg.Error as e:
            out["B"] = f"SQLSTATE {e.sqlstate}"; b.rollback()
    tb = threading.Thread(target=b_side)
    t = time.perf_counter(); tb.start(); time.sleep(0.3)
    try:
        a.commit()
    except psycopg.Error as e:
        out["A"] = f"SQLSTATE {e.sqlstate}"; a.rollback()
    tb.join(); waited = time.perf_counter() - t
    a.close(); b.close()
    return {"variant": f"{level}{' + shared row' if materialize else ''}", "both_saw": seen[0],
            "A": out["A"], "B": out["B"], "final_sum": total(), "b_finished_after_s": waited}


rows = [run("REPEATABLE READ", False), run("SERIALIZABLE", False), run("REPEATABLE READ", True)]
for r in rows:
    print(f"{r['variant']:<30} both read {r['both_saw']}   A: {r['A']:<15} B: {r['B']:<15} final sum {r['final_sum']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"variants": rows, "summary": [{"rr_sum": rows[0]["final_sum"], "ser_sum": rows[1]["final_sum"],
                   "ser_b": rows[1]["B"], "mat_sum": rows[2]["final_sum"], "mat_b": rows[2]["B"]}]}}, f)