#!/usr/bin/env python3
"""Capture harness for the DDIA Ch. 8 in-doubt two-phase-commit exercise.

Drives a REAL prepared (in-doubt) transaction on a REAL PostgreSQL 16 server
and captures, verbatim, what it does to locks and how it survives a full server
restart. The heart of the exercise is phase 1 of two-phase commit:

    Session A: BEGIN;
               UPDATE accounts SET balance = balance - 100 WHERE id = 1;
               PREPARE TRANSACTION 'tx1';

After PREPARE the transaction is neither committed nor rolled back: it is
"prepared" -- durably written to disk, owned by NO session, and STILL HOLDING
the exclusive row lock it took. This harness shows, from real output, that:

    1. pg_prepared_xacts lists tx1 while it sits in doubt.
    2. A second session's UPDATE ... WHERE id = 1 BLOCKS on tx1's lock.
    3. A full `docker restart` of the container does NOT clear it: tx1 is
       STILL listed in pg_prepared_xacts and STILL holds the lock afterwards
       (it was persisted, exactly as 2PC requires).
    4. Only an explicit ROLLBACK PREPARED 'tx1' (or COMMIT PREPARED) frees the
       lock -- at which point the waiting second session immediately proceeds.

This is the CAPTURE / VERIFY path, not the reader path. Readers of the exercise
type the psql sessions by hand on port 5433 with a `docker restart` in between.
This script drives the same sequence -- including the restart via subprocess and
a blocking second session on a thread -- so the printed transcript is real,
verbatim output the write-up quotes.

Because this harness runs `docker restart`, it interrupts its OWN connections.
It is written to expect that: every connection opened before the restart is
abandoned, and it reconnects cleanly afterwards (waiting for pg_isready first).

Idempotent / re-runnable: at start it ROLLBACK PREPAREDs any leftover 'tx1' on
the shared server, then drops and recreates its own database `ex_ch08_2pc`.

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

Substrate (already running; this script restarts ONLY this one container):
    PostgreSQL 16 container `ddia-ch08-2pc` on localhost:5433, user=postgres,
    trust auth, maintenance db `study`, started with max_prepared_transactions=10.
"""

import subprocess
import threading
import time

import psycopg

CONTAINER = "ddia-ch08-2pc"
PORT = 5433
ADMIN_DSN = f"host=localhost port={PORT} user=postgres dbname=study"
DBNAME = "ex_ch08_2pc"
EX_DSN = f"host=localhost port={PORT} user=postgres dbname={DBNAME}"
GID = "tx1"


def clear_leftover_prepared() -> None:
    """ROLLBACK PREPARED any stale gid from a previous interrupted run.

    A prepared xact outlives its database's session but not the cluster, so a
    leftover 'tx1' can even block DROP DATABASE. Resolve it first, in whatever
    database it belongs to."""
    with psycopg.connect(ADMIN_DSN, autocommit=True) as admin:
        rows = admin.execute(
            "SELECT gid, database FROM pg_prepared_xacts WHERE gid = %s", (GID,)
        ).fetchall()
        for gid, db in rows:
            dsn = f"host=localhost port={PORT} user=postgres dbname={db}"
            with psycopg.connect(dsn, autocommit=True) as c:
                c.execute(f"ROLLBACK PREPARED '{gid}'")
                print(f"(cleanup) rolled back leftover prepared xact '{gid}' in db {db}")


def recreate_database() -> None:
    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:
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("CREATE TABLE accounts (id int PRIMARY KEY, balance int)")
        conn.execute("INSERT INTO accounts VALUES (1, 500)")


def prepared_xacts_rows(conn):
    """The real pg_prepared_xacts rows for our gid: (gid, owner, database)."""
    return conn.execute(
        "SELECT gid, owner, database FROM pg_prepared_xacts WHERE gid = %s "
        "ORDER BY gid",
        (GID,),
    ).fetchall()


def prepared_holds_lock(conn) -> bool:
    """True if a prepared transaction (a lock with no live pid) holds a lock
    on the accounts table -- the in-doubt txn's grip, visible in pg_locks."""
    row = conn.execute(
        """
        SELECT count(*) FROM pg_locks l
        WHERE l.relation = 'accounts'::regclass
          AND l.pid IS NULL          -- prepared xacts have no backend pid
          AND l.granted
        """
    ).fetchone()
    return row[0] > 0


def balance(conn) -> int:
    return conn.execute("SELECT balance FROM accounts WHERE id = 1").fetchone()[0]


def restart_container() -> None:
    print(f"--- docker restart {CONTAINER} (full server restart) ---")
    subprocess.run(["docker", "restart", CONTAINER], check=True,
                   stdout=subprocess.DEVNULL)
    # Wait for the server to accept connections again.
    for _ in range(60):
        r = subprocess.run(
            ["docker", "exec", CONTAINER, "pg_isready", "-U", "postgres"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        if r.returncode == 0:
            break
        time.sleep(0.5)
    # pg_isready can flip to accepting a beat before it truly serves queries.
    for _ in range(60):
        try:
            with psycopg.connect(EX_DSN, connect_timeout=2) as c:
                c.execute("SELECT 1")
            break
        except psycopg.OperationalError:
            time.sleep(0.5)
    print("--- server back up ---\n")


class BlockingSession:
    """A second session that runs UPDATE ... WHERE id = 1 in a thread.

    Because tx1 holds the row's exclusive lock, this UPDATE blocks. The thread
    stays alive (blocked) until the lock is released, then completes. `done`
    is set the instant the UPDATE returns."""

    def __init__(self, label):
        self.label = label
        self.done = threading.Event()
        self.error = None
        self.conn = None
        self.thread = threading.Thread(target=self._run)

    def _run(self):
        try:
            self.conn = psycopg.connect(EX_DSN, autocommit=False)
            with self.conn.cursor() as cur:
                cur.execute("UPDATE accounts SET balance = balance + 1 WHERE id = 1")
            self.conn.commit()
        except Exception as e:  # noqa: BLE001 -- report whatever really happened
            self.error = e
        finally:
            self.done.set()

    def start(self):
        self.thread.start()

    def wait_until_blocking_visible(self, admin_conn, timeout=5.0):
        """Poll pg_stat_activity until this session is visibly waiting on a lock.
        Returns the (state, wait_event_type, wait_event) tuple it observed."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            row = admin_conn.execute(
                """
                SELECT state, wait_event_type, wait_event
                FROM pg_stat_activity
                WHERE datname = %s
                  AND query ILIKE 'UPDATE accounts%%'
                  AND wait_event_type = 'Lock'
                """,
                (DBNAME,),
            ).fetchone()
            if row is not None:
                return row
            time.sleep(0.1)
        return None

    def close(self):
        try:
            if self.conn is not None:
                self.conn.close()
        except Exception:  # noqa: BLE001
            pass


def hr(title):
    print(f"\n========== {title} ==========")


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

    print("accounts seeded: id=1, balance=500\n")

    # --- Phase 1: prepare tx1 (2PC phase 1) --------------------------------
    hr("Session A: BEGIN; UPDATE; PREPARE TRANSACTION 'tx1'")
    a = psycopg.connect(EX_DSN, autocommit=False)
    with a.cursor() as cur:
        cur.execute("UPDATE accounts SET balance = balance - 100 WHERE id = 1")
        cur.execute(f"PREPARE TRANSACTION '{GID}'")
    # After PREPARE, session A's transaction is finished; the prepared xact is
    # now detached from this (or any) session. We can drop the connection.
    a.close()
    print("PREPARE TRANSACTION 'tx1' issued -- txn is now PREPARED (in doubt).")
    print("Session A is free; tx1 is owned by no session.\n")

    obs = psycopg.connect(EX_DSN, autocommit=True)
    print("SELECT gid, owner, database FROM pg_prepared_xacts;")
    for r in prepared_xacts_rows(obs):
        print(f"  gid={r[0]}  owner={r[1]}  database={r[2]}")
    print(f"prepared xact holds a lock on accounts: {prepared_holds_lock(obs)}")
    print(f"balance of id=1 (uncommitted -100 is invisible): {balance(obs)}")

    # --- Phase 1b: a second session blocks on tx1's lock -------------------
    hr("Session B: UPDATE accounts WHERE id=1  ->  blocks on tx1")
    b1 = BlockingSession("B(pre-restart)")
    b1.start()
    waited = b1.wait_until_blocking_visible(obs)
    if waited:
        print(f"Session B state={waited[0]}  wait={waited[1]}:{waited[2]}  "
              f"(blocked on the in-doubt txn's lock)")
    still_blocked = not b1.done.wait(timeout=2.0)
    print(f"after 2s more, Session B still blocked: {still_blocked}")
    obs.close()
    # The restart will kill this blocked connection; abandon it.
    b1.close()

    # --- Phase 2: full server restart --------------------------------------
    hr("docker restart -- does a full restart clear the in-doubt txn?")
    restart_container()

    obs2 = psycopg.connect(EX_DSN, autocommit=True)
    print("SELECT gid, owner, database FROM pg_prepared_xacts;  (AFTER restart)")
    rows_after = prepared_xacts_rows(obs2)
    for r in rows_after:
        print(f"  gid={r[0]}  owner={r[1]}  database={r[2]}")
    print(f"tx1 STILL listed after restart: {len(rows_after) > 0}")
    print(f"prepared xact STILL holds lock after restart: {prepared_holds_lock(obs2)}")
    print(f"balance of id=1 after restart: {balance(obs2)}")

    # A fresh second session still blocks -- the lock survived the restart.
    hr("Session B (fresh): UPDATE WHERE id=1  ->  STILL blocks after restart")
    b2 = BlockingSession("B(post-restart)")
    b2.start()
    waited2 = b2.wait_until_blocking_visible(obs2)
    if waited2:
        print(f"Session B state={waited2[0]}  wait={waited2[1]}:{waited2[2]}  "
              f"(still blocked -- the persisted lock is real)")
    still_blocked2 = not b2.done.wait(timeout=2.0)
    print(f"after 2s more, Session B still blocked: {still_blocked2}")

    # --- Phase 3: resolve the in-doubt txn ---------------------------------
    hr("ROLLBACK PREPARED 'tx1' -- resolve the in-doubt txn")
    obs2.execute(f"ROLLBACK PREPARED '{GID}'")
    print("ROLLBACK PREPARED 'tx1' issued.")

    # The still-waiting Session B should now proceed immediately.
    proceeded = b2.done.wait(timeout=5.0)
    b2.thread.join(timeout=5.0)
    print(f"Session B unblocked and completed after the rollback: {proceeded}")
    if b2.error:
        print(f"Session B error: {b2.error!r}")
    b2.close()

    print(f"tx1 still in pg_prepared_xacts: {len(prepared_xacts_rows(obs2)) > 0}")
    print(f"prepared xact still holds lock: {prepared_holds_lock(obs2)}")
    # Balance: rollback undid the -100 (back to 500); B's committed +1 makes 501.
    print(f"balance of id=1 (tx1 rolled back to 500, +1 from Session B): {balance(obs2)}")
    obs2.close()

    hr("summary")
    print("- PREPARE TRANSACTION left tx1 in doubt, holding the row lock.")
    print("- A second session's UPDATE blocked on it indefinitely.")
    print("- A full `docker restart` did NOT clear it: still listed, still locking.")
    print("- Only ROLLBACK PREPARED freed the lock; the waiter then proceeded.")


if __name__ == "__main__":
    main()