#!/usr/bin/env python3
"""Capture harness for the DDIA Ch. 8 phantom / write-skew (double-booking) exercise.

Reproduces the meeting-room double-booking anomaly (DDIA 2e Example 8-2) on a
REAL PostgreSQL 16 server. Two transactions each check that a room is free for a
time slot -- SELECT count(*) of overlapping bookings, both see 0 -- then each
INSERT a booking for the SAME room and slot, and both commit. The conflict is
over a row that does NOT exist yet: a *phantom*. Four scenarios:

  1. REPEATABLE READ, plain check         -> BOTH commit  -> DOUBLE-BOOKED (2 rows)
  2. REPEATABLE READ, SELECT ... FOR UPDATE-> BOTH commit  -> STILL double-booked
                                             (FOR UPDATE locks existing rows; there
                                              are none, so it locks nothing)
  3. SERIALIZABLE                          -> one commits, the other aborts with a
                                             REAL 40001 serialization failure (SSI
                                             tracks the predicate you queried)
  4. EXCLUDE USING gist constraint         -> one commits, the other's INSERT is
                                             rejected with a REAL 23P01 exclusion
                                             violation (a physical stand-in for the
                                             predicate)

This is the CAPTURE / VERIFY path, not the reader path. Readers of the exercise
type two `psql` sessions by hand (Session A / Session B); this script drives the
same two sessions programmatically, holding both connections open at once and
interleaving them, so the printed transcript is real, verbatim output that the
write-up quotes.

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

Run:
    uv run --python 3.12 --with "psycopg[binary]" python3 code/phantom_write_skew.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 psycopg
from psycopg import errors

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

# The slot both transactions try to book: room 1, 10:00-11:00.
ROOM, START_H, END_H = 1, 10, 11

# The overlap predicate, written the way a reader would type it in psql:
# an existing booking [s, e) overlaps the requested [10, 11) iff s < 11 AND e > 10.
OVERLAP = f"room_id = {ROOM} AND start_hour < {END_H} AND end_hour > {START_H}"


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 create_bookings(with_exclusion: bool = False) -> None:
    """Create an EMPTY bookings table (no overlapping booking exists yet).

    If with_exclusion, add an EXCLUDE USING gist constraint that forbids two
    bookings of the same room whose hour-ranges overlap (needs btree_gist for
    the integer equality half of the constraint).
    """
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        conn.execute("DROP TABLE IF EXISTS bookings")
        conn.execute(
            "CREATE TABLE bookings ("
            "  id serial PRIMARY KEY,"
            "  room_id int NOT NULL,"
            "  start_hour int NOT NULL,"
            "  end_hour int NOT NULL)"
        )
        if with_exclusion:
            conn.execute("CREATE EXTENSION IF NOT EXISTS btree_gist")
            conn.execute(
                "ALTER TABLE bookings ADD CONSTRAINT no_double_booking "
                "EXCLUDE USING gist ("
                "  room_id WITH =,"
                "  int4range(start_hour, end_hour) WITH &&)"
            )


def count_overlap(cur, for_update: bool = False) -> int:
    """How many existing bookings overlap room 1, 10:00-11:00?

    With for_update=True this locks the matching rows FOR UPDATE -- the lost-update
    fix. On an empty overlap set it locks nothing, which is the whole point.
    Note: FOR UPDATE cannot be combined with count(*), so we select the rows and
    count them client-side, exactly as the reader would eyeball the row count.
    """
    sql = f"SELECT id FROM bookings WHERE {OVERLAP}"
    if for_update:
        sql += " FOR UPDATE"
    cur.execute(sql)
    return len(cur.fetchall())


def insert_booking(cur) -> None:
    cur.execute(
        "INSERT INTO bookings (room_id, start_hour, end_hour) VALUES (%s, %s, %s)",
        (ROOM, START_H, END_H),
    )


def total_bookings() -> int:
    with psycopg.connect(EX_DSN, autocommit=True) as conn:
        return conn.execute(
            f"SELECT count(*) FROM bookings WHERE {OVERLAP}"
        ).fetchone()[0]


def run_interleaving(level: str, for_update: bool) -> None:
    """Drive the classic write-skew interleaving with two live connections.

    A: check (0) ; B: check (0) ; A: insert ; B: insert ; A: commit ; B: commit.
    Both connections are open at the same time so the checks genuinely precede
    both inserts -- the interleaving a reader builds by hand across two psql tabs.
    """
    fu = "  (SELECT ... FOR UPDATE)" if for_update else ""
    print(f"=== isolation level: {level}{fu} ===")

    sa = psycopg.connect(EX_DSN, autocommit=False)
    sb = psycopg.connect(EX_DSN, autocommit=False)
    try:
        sa.execute(f"SET TRANSACTION ISOLATION LEVEL {level}")
        sb.execute(f"SET TRANSACTION ISOLATION LEVEL {level}")

        with sa.cursor() as ca, sb.cursor() as cb:
            # Both check first -- neither has inserted yet, so both see 0.
            a_seen = count_overlap(ca, for_update)
            print(f"Session A  check: overlapping bookings for room 1, 10-11  ->  {a_seen}")
            b_seen = count_overlap(cb, for_update)
            print(f"Session B  check: overlapping bookings for room 1, 10-11  ->  {b_seen}")

            # Both believe the room is free; both insert.
            insert_booking(ca)
            print("Session A  INSERT booking (room 1, 10-11)")
            insert_booking(cb)
            print("Session B  INSERT booking (room 1, 10-11)")

            # Commit A, then B.
            try:
                sa.commit()
                print("Session A  COMMIT  ->  OK")
            except psycopg.Error as e:
                print(f"Session A  COMMIT  ->  {e.sqlstate}  {e.diag.message_primary}")
            try:
                sb.commit()
                print("Session B  COMMIT  ->  OK")
            except psycopg.Error as e:
                print(f"Session B  COMMIT  ->  ERROR  SQLSTATE {e.sqlstate}")
                print(f"           {e.diag.message_primary}")
    finally:
        sa.close()
        sb.close()

    print(f"result: {total_bookings()} booking(s) exist for room 1, 10-11\n")


def run_exclusion() -> None:
    """Same interleaving, but the table has an EXCLUDE USING gist constraint.

    The constraint is a physical stand-in for the overlap predicate. B's INSERT
    blocks on A's uncommitted row; once A commits, B's INSERT is rejected with a
    real 23P01 exclusion_violation. Runs at READ COMMITTED to show the constraint,
    not the isolation level, is doing the work.
    """
    print("=== EXCLUDE USING gist constraint (READ COMMITTED) ===")

    sa = psycopg.connect(EX_DSN, autocommit=False)
    sb = psycopg.connect(EX_DSN, autocommit=False)
    try:
        with sa.cursor() as ca, sb.cursor() as cb:
            a_seen = count_overlap(ca)
            print(f"Session A  check: overlapping bookings for room 1, 10-11  ->  {a_seen}")
            b_seen = count_overlap(cb)
            print(f"Session B  check: overlapping bookings for room 1, 10-11  ->  {b_seen}")

            insert_booking(ca)
            print("Session A  INSERT booking (room 1, 10-11)")
            sa.commit()
            print("Session A  COMMIT  ->  OK")

            # B's INSERT now collides with A's committed row via the constraint.
            try:
                insert_booking(cb)
                sb.commit()
                print("Session B  COMMIT  ->  OK")
            except psycopg.Error as e:
                print(f"Session B  INSERT  ->  ERROR  SQLSTATE {e.sqlstate}")
                print(f"           {e.diag.message_primary}")
                if e.diag.message_detail:
                    print(f"           DETAIL:  {e.diag.message_detail}")
                sb.rollback()
    finally:
        sa.close()
        sb.close()

    print(f"result: {total_bookings()} booking(s) exist for room 1, 10-11\n")


def main() -> None:
    recreate_database()

    print("An empty bookings table. Two sessions each try to book room 1, "
          "10:00-11:00.\nThe invariant: at most ONE booking may exist for that slot.\n")

    # 1. REPEATABLE READ, plain check -> double-booked.
    create_bookings()
    run_interleaving("REPEATABLE READ", for_update=False)

    # 2. REPEATABLE READ + SELECT ... FOR UPDATE -> STILL double-booked.
    create_bookings()
    run_interleaving("REPEATABLE READ", for_update=True)

    # 3. SERIALIZABLE -> one aborts with a real 40001.
    create_bookings()
    run_interleaving("SERIALIZABLE", for_update=False)

    # 4. EXCLUDE USING gist -> one INSERT rejected with a real 23P01.
    create_bookings(with_exclusion=True)
    run_exclusion()


if __name__ == "__main__":
    main()