#!/usr/bin/env python3
"""Capture harness for the DDIA Ch. 8 deadlock-detection exercise.

Triggers a REAL deadlock on a REAL PostgreSQL 16 server and captures the
verbatim error the database returns when its deadlock detector fires. Two
transactions lock the same two rows in the OPPOSITE order:

    Thread A: BEGIN; UPDATE rows SET v=v+1 WHERE id=1;  (locks row 1)
              ... then UPDATE rows SET v=v+1 WHERE id=2;  (wants row 2)
    Thread B: BEGIN; UPDATE rows SET v=v+1 WHERE id=2;  (locks row 2)
              ... then UPDATE rows SET v=v+1 WHERE id=1;  (wants row 1)

Each holds one row and waits for the other -> a cycle in the wait-for graph.
Under two-phase locking (which Postgres uses for row write-locks) neither can
proceed. After `deadlock_timeout` (default ~1 s) Postgres's deadlock detector
notices the cycle and aborts ONE victim with SQLSTATE 40P01 "deadlock
detected"; the survivor proceeds and commits.

This is the CAPTURE / VERIFY path, not the reader path. Readers of the exercise
type two `psql` sessions by hand (Session A / Session B) and interleave the two
UPDATEs. This script drives the same two transactions from two threads so the
printed transcript -- including the exact 40P01 error and its DETAIL -- is real,
verbatim output the write-up quotes.

It is deterministic in the DETECTION (a victim is always chosen and the
survivor always commits) but NOT in the victim's identity: which thread loses
the race to grab the second lock, and therefore which one Postgres aborts, can
vary run to run. The script reports whichever it observed. It is idempotent:
drops and recreates its own database `ex_ch08_deadlock` on the shared,
already-running server, and closes every connection before dropping.

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

Substrate (already running; this script does NOT manage containers):
    PostgreSQL 16 container `ddia-ch08-pg` on localhost:5432, user=postgres,
    trust auth, maintenance db `study`.
"""

import threading
import time

import psycopg

ADMIN_DSN = "host=localhost port=5432 user=postgres dbname=study"
DBNAME = "ex_ch08_deadlock"
EX_DSN = f"host=localhost port=5432 user=postgres dbname={DBNAME}"


def recreate_database() -> None:
    """Drop and recreate our own database via an autocommit admin connection."""
    with psycopg.connect(ADMIN_DSN, autocommit=True) as admin:
        admin.execute(f"DROP DATABASE IF EXISTS {DBNAME} WITH (FORCE)")
        admin.execute(f"CREATE DATABASE {DBNAME}")


def seed() -> None:
    """Create the rows table and seed two rows at v=0."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("DROP TABLE IF EXISTS rows")
        conn.execute("CREATE TABLE rows (id int PRIMARY KEY, v int)")
        conn.execute("INSERT INTO rows VALUES (1, 0), (2, 0)")


def run_thread(name, first_id, second_id, gate, results):
    """One transaction: lock `first_id`, wait at the gate, then want `second_id`.

    The barrier `gate` guarantees BOTH transactions have taken their first row
    lock before EITHER reaches for its second -- that is what forms the cycle.
    Whichever transaction Postgres picks as the victim raises DeadlockDetected
    (SQLSTATE 40P01) on its second UPDATE; the survivor's second UPDATE succeeds
    and it commits.
    """
    conn = psycopg.connect(EX_DSN, autocommit=False)
    try:
        with conn.cursor() as cur:
            # First UPDATE: take a row-level write lock on `first_id`.
            cur.execute("UPDATE rows SET v = v + 1 WHERE id = %s", (first_id,))
            print(f"{name}  UPDATE id={first_id}  ->  locked row {first_id}")

            # Wait until the OTHER thread has also taken its first lock.
            gate.wait()

            # Second UPDATE: reach for the row the other thread holds.
            print(f"{name}  UPDATE id={second_id}  ->  waiting for row {second_id}...")
            cur.execute("UPDATE rows SET v = v + 1 WHERE id = %s", (second_id,))

        conn.commit()
        print(f"{name}  COMMIT  ->  survivor: both updates applied")
        results[name] = ("survivor", None)
    except psycopg.errors.DeadlockDetected as e:
        conn.rollback()
        print(f"{name}  ABORTED by deadlock detector (SQLSTATE {e.sqlstate})")
        results[name] = ("victim", str(e).rstrip())
    finally:
        conn.close()


def main() -> None:
    recreate_database()
    seed()

    print("Two rows, (1, v=0) and (2, v=0). Two transactions lock them in")
    print("opposite order: A takes row 1 then wants row 2; B takes row 2 then")
    print("wants row 1. Each waits on the other -- a cycle.\n")

    gate = threading.Barrier(2)
    results = {}

    # Thread A: row 1 first, then row 2. Thread B: row 2 first, then row 1.
    a = threading.Thread(target=run_thread, args=("A", 1, 2, gate, results))
    b = threading.Thread(target=run_thread, args=("B", 2, 1, gate, results))

    t0 = time.monotonic()
    a.start()
    b.start()
    a.join()
    b.join()
    elapsed = time.monotonic() - t0

    print()
    victim = next(n for n, (role, _) in results.items() if role == "victim")
    survivor = next(n for n, (role, _) in results.items() if role == "survivor")
    err = results[victim][1]

    print("=== the 40P01 error the victim received (verbatim) ===")
    print(err)
    print()

    # Show the survivor's committed effect.
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        rows = conn.execute("SELECT id, v FROM rows ORDER BY id").fetchall()

    print("=== summary ===")
    print(f"victim   : Thread {victim}  (aborted, SQLSTATE 40P01, rolled back)")
    print(f"survivor : Thread {survivor}  (committed)")
    print(f"final rows after survivor commit: {rows}")
    print(f"time from both-locked to detection + abort: ~{elapsed:.2f}s "
          f"(~deadlock_timeout, default 1s -- the detector firing, not a benchmark)")


if __name__ == "__main__":
    main()