#!/usr/bin/env python3
"""Capture harness for DDIA Ch. 8 exercise: "The lost update, and three fixes".

Reproduces, against a REAL PostgreSQL 16, the classic lost update anomaly and
the three fixes the book describes:

    Scenario 0 (the bug)  two app-level read-modify-write increments under
                          READ COMMITTED -> final value 43, not 44. One
                          increment is silently lost.
    Fix 1  atomic in-DB   UPDATE ... SET value = value + 1  -> 44 (safe even
                          at READ COMMITTED; no app-side read/write gap).
    Fix 2  REPEATABLE      the app-level read-then-write, but Postgres DETECTS
           READ            the write-write conflict and ABORTS the loser with
                          SQLSTATE 40001 "could not serialize access due to
                          concurrent update". You must catch it and retry.
    Fix 3  SELECT ...      lock the row on read; the second session blocks
           FOR UPDATE      until the first commits, then re-reads -> 44.

Two sessions (A and B) run against one database. The FOR UPDATE scenario uses a
real background thread so session B genuinely *blocks* on A's row lock; every
other scenario is strictly sequential. All four final values and the 40001
error text are captured from the live server -- nothing here is invented.

Deterministic and idempotent: it drops+recreates the database
`ex_ch08_lostupdate` every run, reseeds value=42 before each scenario, and
closes all connections before dropping. This is the maintainer's verification
path -- NOT the reader's path. The reader reproduces the same thing by hand in
two psql sessions (see lost-update.md, Steps).

Run:
    uv run --python 3.12 --with "psycopg[binary]" python3 code/lost_update.py
"""

import threading
import time

import psycopg

ADMIN_DSN = "host=localhost port=5432 user=postgres dbname=study"
DB = "ex_ch08_lostupdate"
DSN = f"host=localhost port=5432 user=postgres dbname={DB}"


def rule():
    print("=" * 70)


def reset_database():
    """Drop and recreate ex_ch08_lostupdate via an autocommit admin connection."""
    with psycopg.connect(ADMIN_DSN, autocommit=True) as admin:
        admin.execute(f"DROP DATABASE IF EXISTS {DB} WITH (FORCE)")
        admin.execute(f"CREATE DATABASE {DB}")


def create_table():
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("CREATE TABLE counters (id int PRIMARY KEY, value int)")


def reseed():
    """Set the one counter row back to (1, 42), committed, before a scenario."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("DELETE FROM counters")
        conn.execute("INSERT INTO counters VALUES (1, 42)")


def final_value():
    with psycopg.connect(DSN, autocommit=True) as conn:
        return conn.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]


# --------------------------------------------------------------------------
# Scenario 0 -- the bug: two app-level read-modify-writes, READ COMMITTED.
# --------------------------------------------------------------------------
def scenario_bug():
    reseed()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)

    rule()
    print("SCENARIO 0  the lost update  (READ COMMITTED, app-level read-modify-write)")
    rule()
    print("  counter seeded:  id=1  value=42")
    print()

    # A reads.
    a.execute("BEGIN")  # READ COMMITTED is the default
    a_read = a.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]
    print("SESSION A:  BEGIN;")
    print(f"SESSION A:  SELECT value FROM counters WHERE id = 1;   -- {a_read}")

    # B reads the SAME committed value.
    b.execute("BEGIN")
    b_read = b.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]
    print(f"SESSION B:  BEGIN;")
    print(f"SESSION B:  SELECT value FROM counters WHERE id = 1;   -- {b_read}")

    # A computes value+1 IN APP CODE and writes it back, then commits.
    a_new = a_read + 1
    a.execute("UPDATE counters SET value = %s WHERE id = 1", (a_new,))
    a.commit()
    print(f"SESSION A:  UPDATE counters SET value = {a_new} WHERE id = 1;   -- computed {a_read}+1 in app")
    print("SESSION A:  COMMIT;")

    # B does the same, from ITS stale read of 42.
    b_new = b_read + 1
    b.execute("UPDATE counters SET value = %s WHERE id = 1", (b_new,))
    b.commit()
    print(f"SESSION B:  UPDATE counters SET value = {b_new} WHERE id = 1;   -- computed {b_read}+1 in app")
    print("SESSION B:  COMMIT;")

    a.close()
    b.close()
    fv = final_value()
    print()
    print(f"  final value = {fv}   (expected 44 if both increments survived)")
    return fv


# --------------------------------------------------------------------------
# Fix 1 -- atomic in-DB increment. Safe even at READ COMMITTED.
# --------------------------------------------------------------------------
def scenario_atomic():
    reseed()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)

    rule()
    print("FIX 1  atomic in-DB increment  (UPDATE ... SET value = value + 1)")
    rule()
    print("  counter seeded:  id=1  value=42")
    print()

    a.execute("BEGIN")
    a.execute("UPDATE counters SET value = value + 1 WHERE id = 1")
    a.commit()
    print("SESSION A:  BEGIN;")
    print("SESSION A:  UPDATE counters SET value = value + 1 WHERE id = 1;")
    print("SESSION A:  COMMIT;")

    b.execute("BEGIN")
    b.execute("UPDATE counters SET value = value + 1 WHERE id = 1")
    b.commit()
    print("SESSION B:  BEGIN;")
    print("SESSION B:  UPDATE counters SET value = value + 1 WHERE id = 1;")
    print("SESSION B:  COMMIT;")

    a.close()
    b.close()
    fv = final_value()
    print()
    print(f"  final value = {fv}   (no app-side read; each UPDATE reads+writes atomically)")
    return fv


# --------------------------------------------------------------------------
# Fix 2 -- REPEATABLE READ: Postgres DETECTS the conflict and aborts the loser.
# --------------------------------------------------------------------------
def scenario_repeatable_read():
    reseed()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)

    rule()
    print("FIX 2  REPEATABLE READ  (snapshot isolation detects the write-write conflict)")
    rule()
    print("  counter seeded:  id=1  value=42")
    print()

    a.execute("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    a_read = a.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]
    print("SESSION A:  BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;")
    print(f"SESSION A:  SELECT value FROM counters WHERE id = 1;   -- {a_read}")

    b.execute("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    b_read = b.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]
    print("SESSION B:  BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;")
    print(f"SESSION B:  SELECT value FROM counters WHERE id = 1;   -- {b_read}")

    a_new = a_read + 1
    a.execute("UPDATE counters SET value = %s WHERE id = 1", (a_new,))
    a.commit()
    print(f"SESSION A:  UPDATE counters SET value = {a_new} WHERE id = 1;   -- computed {a_read}+1 in app")
    print("SESSION A:  COMMIT;")

    # B now tries the same write from its own snapshot -> serialization failure.
    b_new = b_read + 1
    sqlstate = None
    errmsg = None
    try:
        b.execute("UPDATE counters SET value = %s WHERE id = 1", (b_new,))
        b.commit()
        print(f"SESSION B:  UPDATE counters SET value = {b_new} WHERE id = 1;   -- UNEXPECTED: no error")
    except psycopg.errors.SerializationFailure as e:
        sqlstate = e.sqlstate
        errmsg = str(e).strip()
        b.rollback()
        print(f"SESSION B:  UPDATE counters SET value = {b_new} WHERE id = 1;")
        print(f"SESSION B:  -- ERROR  SQLSTATE {sqlstate}")
        print(f"SESSION B:  -- {errmsg}")
        print("SESSION B:  ROLLBACK;  (you must catch this and RETRY)")

    # The retry: a fresh REPEATABLE READ txn sees A's committed 43 and adds 1.
    b.execute("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    b_retry_read = b.execute("SELECT value FROM counters WHERE id = 1").fetchone()[0]
    b_retry_new = b_retry_read + 1
    b.execute("UPDATE counters SET value = %s WHERE id = 1", (b_retry_new,))
    b.commit()
    print("SESSION B:  BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;   -- retry")
    print(f"SESSION B:  SELECT value FROM counters WHERE id = 1;   -- {b_retry_read}")
    print(f"SESSION B:  UPDATE counters SET value = {b_retry_new} WHERE id = 1;   -- computed {b_retry_read}+1")
    print("SESSION B:  COMMIT;")

    a.close()
    b.close()
    fv = final_value()
    print()
    print(f"  final value = {fv}   (loser aborted, retried on the fresh value)")
    return fv, sqlstate, errmsg


# --------------------------------------------------------------------------
# Fix 3 -- SELECT ... FOR UPDATE: lock on read, the second session blocks.
# --------------------------------------------------------------------------
def scenario_for_update():
    reseed()
    a = psycopg.connect(DSN)
    b = psycopg.connect(DSN)

    rule()
    print("FIX 3  SELECT ... FOR UPDATE  (explicit row lock serializes the two sessions)")
    rule()
    print("  counter seeded:  id=1  value=42")
    print()

    # A locks the row on read.
    a.execute("BEGIN")
    a_read = a.execute(
        "SELECT value FROM counters WHERE id = 1 FOR UPDATE"
    ).fetchone()[0]
    print("SESSION A:  BEGIN;")
    print(f"SESSION A:  SELECT value FROM counters WHERE id = 1 FOR UPDATE;   -- {a_read}  (row now locked)")

    # B's FOR UPDATE will block on A's lock; run it in a thread and capture
    # the value it eventually reads once A releases.
    b_result = {}

    def b_worker():
        b.execute("BEGIN")
        t0 = time.monotonic()
        val = b.execute(
            "SELECT value FROM counters WHERE id = 1 FOR UPDATE"
        ).fetchone()[0]
        b_result["waited"] = time.monotonic() - t0
        b_result["read"] = val
        new = val + 1
        b.execute("UPDATE counters SET value = %s WHERE id = 1", (new,))
        b.commit()
        b_result["new"] = new

    t = threading.Thread(target=b_worker)
    t.start()
    time.sleep(0.5)  # let B reach its FOR UPDATE and block on A's lock
    print("SESSION B:  BEGIN;")
    print("SESSION B:  SELECT value FROM counters WHERE id = 1 FOR UPDATE;   -- BLOCKS (waiting for A)")

    # A finishes; releasing the lock unblocks B.
    a_new = a_read + 1
    a.execute("UPDATE counters SET value = %s WHERE id = 1", (a_new,))
    a.commit()
    print(f"SESSION A:  UPDATE counters SET value = {a_new} WHERE id = 1;   -- computed {a_read}+1")
    print("SESSION A:  COMMIT;   -- releases the lock")

    t.join()
    print(f"SESSION B:  -- unblocks after ~{b_result['waited']:.1f}s, FOR UPDATE returns {b_result['read']}")
    print(f"SESSION B:  UPDATE counters SET value = {b_result['new']} WHERE id = 1;   -- computed {b_result['read']}+1")
    print("SESSION B:  COMMIT;")

    a.close()
    b.close()
    fv = final_value()
    print()
    print(f"  final value = {fv}   (the lock forced B to read A's committed value first)")
    return fv


def main():
    reset_database()
    create_table()

    bug = scenario_bug()
    print()
    atomic = scenario_atomic()
    print()
    rr, sqlstate, errmsg = scenario_repeatable_read()
    print()
    foru = scenario_for_update()

    print()
    rule()
    print("SUMMARY")
    rule()
    print(f"  Scenario 0  app read-modify-write, READ COMMITTED : {bug}   (lost update)")
    print(f"  Fix 1       UPDATE value = value + 1              : {atomic}")
    print(f"  Fix 2       REPEATABLE READ + retry               : {rr}   (loser SQLSTATE {sqlstate})")
    print(f"  Fix 3       SELECT ... FOR UPDATE                 : {foru}")
    print()
    print(f"  Fix 2 abort message (verbatim): {errmsg}")
    verdict = (
        "as expected: 43 is the lost update; all three fixes reach 44"
        if bug == 43 and atomic == 44 and rr == 44 and foru == 44 and sqlstate == "40001"
        else "UNEXPECTED -- reality diverged from the exercise's aha"
    )
    print(f"  verdict: {verdict}")
    rule()


if __name__ == "__main__":
    main()