# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Reproduce the lost update at Read Committed, then each fix.

Two sessions on one account row in PostgreSQL. Session A reads the balance,
session B adds 50 and commits, session A writes a value computed from what it
read. Then the same interleaving with an atomic update, a version check,
SELECT ... FOR UPDATE, and Repeatable Read.
Needs the PostgreSQL service from the compose file (port 55460).
"""
import argparse, json, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55460/postgres", help="Connection string.")
p.add_argument("--hold", type=float, default=1.0, help="Seconds session A holds its row lock in the FOR UPDATE variant.")
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 accounts")
    admin.execute("CREATE TABLE accounts(id int PRIMARY KEY, balance int, version int DEFAULT 0)")
    admin.execute("INSERT INTO accounts VALUES (7, 100, 0)")
def balance():
    return admin.execute("SELECT balance FROM accounts WHERE id = 7").fetchone()[0]
def session_b():
    with psycopg.connect(args.dsn) as b:              # BEGIN; UPDATE ... + 50; COMMIT (every writer bumps the version)
        b.execute("UPDATE accounts SET balance = balance + 50, version = version + 1 WHERE id = 7")

rows = []
def record(name, note):
    rows.append({"variant": name, "final_balance": balance(), "what_happened": note})
    print(f"{name:<34} final balance {rows[-1]['final_balance']:>4}   {note}")

# 1. Read Committed: A reads 100, B adds 50, A writes 200 -> B's 50 is lost
reset()
with psycopg.connect(args.dsn) as a:
    seen = a.execute("SELECT balance FROM accounts WHERE id = 7").fetchone()[0]
    session_b()
    a.execute("UPDATE accounts SET balance = %s WHERE id = 7", (seen + 100,))
record("Read Committed, write what you read", f"A read {seen} and wrote {seen + 100}; B's +50 is lost")

# 2. Fix: an atomic update computed by the database
reset()
with psycopg.connect(args.dsn) as a:
    a.execute("SELECT balance FROM accounts WHERE id = 7")
    session_b()
    a.execute("UPDATE accounts SET balance = balance + 100 WHERE id = 7")
record("Fix 1: SET balance = balance + 100", "both changes kept")

# 3. Fix: optimistic version check, retry on 0 rows
reset()
retries = 0
with psycopg.connect(args.dsn) as a:
    seen, ver = a.execute("SELECT balance, version FROM accounts WHERE id = 7").fetchone()
    session_b()
    while True:
        cur = a.execute("UPDATE accounts SET balance = %s, version = version + 1 WHERE id = 7 AND version = %s",
                        (seen + 100, ver))
        if cur.rowcount == 1: a.commit(); break
        a.rollback(); retries += 1                       # 0 rows: someone changed it; re-read and retry
        seen, ver = a.execute("SELECT balance, version FROM accounts WHERE id = 7").fetchone()
record("Fix 2: WHERE version = <read>", f"first attempt updated 0 rows; {retries} retry, re-read, then wrote")

# 4. Fix: SELECT ... FOR UPDATE, B blocks until A commits
reset()
waited = {}
with psycopg.connect(args.dsn) as a:
    seen = a.execute("SELECT balance FROM accounts WHERE id = 7 FOR UPDATE").fetchone()[0]
    def b():
        t = time.perf_counter(); session_b(); waited["s"] = time.perf_counter() - t
    tb = threading.Thread(target=b); tb.start()
    time.sleep(args.hold)
    a.execute("UPDATE accounts SET balance = %s WHERE id = 7", (seen + 100,))
tb.join()
record("Fix 3: SELECT ... FOR UPDATE", f"B blocked {waited['s']:.2f} s until A committed, then added 50")

# 5. Repeatable Read: A's write fails with a serialization error instead of losing B's update
reset()
error = ""
with psycopg.connect(args.dsn) as a:
    a.execute("BEGIN ISOLATION LEVEL REPEATABLE READ")
    seen = a.execute("SELECT balance FROM accounts WHERE id = 7").fetchone()[0]
    session_b()
    try:
        a.execute("UPDATE accounts SET balance = %s WHERE id = 7", (seen + 100,))
    except psycopg.errors.SerializationFailure as e:
        error = f"SQLSTATE {e.sqlstate}: {str(e).splitlines()[0]}"
        a.rollback()
record("Repeatable Read, write what you read", error or "no error")
ver = admin.execute("SHOW server_version").fetchone()[0]

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"variants": rows, "summary": [{
            "lost": rows[0]["final_balance"], "atomic": rows[1]["final_balance"], "version": rows[2]["final_balance"],
            "for_update": rows[3]["final_balance"], "blocked_s": waited["s"], "rr": rows[4]["final_balance"],
            "rr_error": error, "server": ver}]}}, f)