#!/usr/bin/env python3
"""Capture harness for the DDIA Ch. 8 "MVCC never blocks readers/writers" exercise.

Demonstrates, against a REAL PostgreSQL 16 server, the multiversion-concurrency-
control guarantee the book states plainly: "readers never block writers, and
writers never block readers." Then it flips ONE knob -- an explicit `SELECT ...
FOR UPDATE` row lock -- to show that when you opt back into locking, a writer
DOES block.

Two cases, same table `items(id, v)` seeded (1, 100):

  Case 1 -- plain snapshot read (non-blocking):
    Session A opens a REPEATABLE READ transaction and SELECTs row 1 (fixing its
    snapshot at v=100). Session B then UPDATEs the same row to 200 and COMMITs.
    B's UPDATE returns immediately -- it does not wait for A. A, still inside its
    transaction, re-SELECTs and STILL sees 100 (its snapshot's version). Neither
    blocked the other.

  Case 2 -- SELECT ... FOR UPDATE (blocking):
    Session A opens a transaction and SELECTs row 1 FOR UPDATE, taking an
    explicit row lock. Session B (a background thread) tries to UPDATE the same
    row. Now B BLOCKS -- it cannot proceed until A releases the lock. A holds the
    lock for a fixed HOLD seconds, then COMMITs; only then does B's UPDATE return.
    B's measured latency is therefore > the hold time.

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 watch the same
behavior; this script drives the two sessions programmatically -- Case 2 needs
real threads and wall-clock timing to *prove* the block -- so the printed
transcript is real, verbatim output the write-up quotes.

It is deterministic and idempotent: it drops and recreates its own database
`ex_ch08_mvcc` on the shared, already-running server, seeds it, runs both cases,
and closes every connection before dropping.

Run:
    uv run --python 3.12 --with "psycopg[binary]" python3 code/mvcc_no_blocking.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_mvcc"
EX_DSN = f"host=localhost port=5432 user=postgres dbname={DBNAME}"

# How long Session A holds the FOR UPDATE row lock in Case 2 before committing.
# The block Session B experiences is measured against this.
HOLD = 1.0


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 items(id, v) and seed one row (1, 100)."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("DROP TABLE IF EXISTS items")
        conn.execute("CREATE TABLE items (id int PRIMARY KEY, v int)")
        conn.execute("INSERT INTO items VALUES (1, 100)")


def reset_row() -> None:
    """Return row 1 to v=100 between cases."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("UPDATE items SET v = 100 WHERE id = 1")


def committed_value() -> int:
    """Read the current committed value of row 1 on a fresh connection."""
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        row = conn.execute("SELECT v FROM items WHERE id = 1").fetchone()
        return row[0]


def case1_plain_read() -> None:
    """Case 1: a plain snapshot read never blocks a concurrent writer."""
    print("=== Case 1: plain read (REPEATABLE READ snapshot) -- non-blocking ===")

    # Session A: open a REPEATABLE READ transaction and fix its snapshot with a
    # plain SELECT of row 1.
    sa = psycopg.connect(EX_DSN, autocommit=False)
    sa.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    a_first = sa.execute("SELECT v FROM items WHERE id = 1").fetchone()[0]
    print(f"Session A  BEGIN (REPEATABLE READ); SELECT v WHERE id=1  ->  {a_first}")

    # Session B: UPDATE the same row and COMMIT. Time how long the UPDATE takes.
    with psycopg.connect(EX_DSN, autocommit=False) as sb:
        t0 = time.perf_counter()
        sb.execute("UPDATE items SET v = 200 WHERE id = 1")
        sb.commit()
        b_latency = time.perf_counter() - t0
    print(f"Session B  UPDATE v=200; COMMIT  ->  returned in {b_latency*1000:.1f} ms "
          f"(did NOT wait for A)")

    # Session A: still inside its transaction, re-SELECT. It sees its snapshot.
    a_second = sa.execute("SELECT v FROM items WHERE id = 1").fetchone()[0]
    print(f"Session A  re-SELECT v WHERE id=1 (same txn)  ->  {a_second} "
          f"(snapshot: still the old version)")
    sa.rollback()
    sa.close()

    print(f"committed value now = {committed_value()} (B's write is durable)")
    print()


def case2_for_update() -> None:
    """Case 2: SELECT ... FOR UPDATE takes a row lock, so the writer blocks."""
    print("=== Case 2: SELECT ... FOR UPDATE -- explicit row lock, writer blocks ===")

    # Session A: open a transaction and lock row 1 with FOR UPDATE.
    sa = psycopg.connect(EX_DSN, autocommit=False)
    sa.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ")
    a_first = sa.execute("SELECT v FROM items WHERE id = 1 FOR UPDATE").fetchone()[0]
    print(f"Session A  BEGIN; SELECT v WHERE id=1 FOR UPDATE  ->  {a_first} "
          f"(row now locked)")

    result = {}
    started = threading.Event()

    def session_b() -> None:
        with psycopg.connect(EX_DSN, autocommit=False) as sb:
            started.set()  # signal we're about to issue the blocking UPDATE
            t0 = time.perf_counter()
            sb.execute("UPDATE items SET v = 300 WHERE id = 1")  # BLOCKS here
            sb.commit()
            result["latency"] = time.perf_counter() - t0

    b = threading.Thread(target=session_b)
    b.start()

    # Wait until B's thread is at the UPDATE, then hold the lock for HOLD seconds.
    started.wait()
    time.sleep(HOLD)
    print(f"Session B  UPDATE v=300  ->  BLOCKED "
          f"(thread alive = {b.is_alive()} while A holds the lock)")

    # Release: A commits, which drops the row lock; B unblocks.
    sa.commit()
    sa.close()
    print(f"Session A  COMMIT  ->  lock released after ~{HOLD:.1f}s hold")

    b.join()
    b_latency = result["latency"]
    print(f"Session B  UPDATE v=300; COMMIT  ->  finally returned in "
          f"{b_latency*1000:.0f} ms (waited > the {HOLD:.1f}s hold)")

    print(f"committed value now = {committed_value()} (B's write applied after A)")
    print()


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

    print("One row: items(id=1, v=100). Two concurrent sessions, A and B.\n")

    reset_row()
    case1_plain_read()

    reset_row()
    case2_for_update()

    print("=== summary ===")
    print("Case 1 (plain read):     B's UPDATE returns immediately; "
          "A keeps reading its snapshot -> readers and writers never block.")
    print("Case 2 (SELECT FOR UPDATE): B's UPDATE blocks until A commits/rolls back "
          "-> an explicit row lock opts back into contention.")


if __name__ == "__main__":
    main()