# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""What query-based CDC misses.

A poller selecting updated_at > :last_seen every second, while another
session (1) updates one row three times between two polls, (2) deletes a
row, and (3) sets updated_at = now() in a transaction that commits 10 s
later. Then the same actions read from a logical replication slot.
Needs the PostgreSQL service from the compose file (port 55480).
"""
import argparse, json, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55480/postgres", help="Connection string.")
p.add_argument("--poll-every", type=float, default=1.0, help="Seconds between polls.")
p.add_argument("--slow-commit", type=float, default=10.0, help="Seconds the slow transaction stays open.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

admin = psycopg.connect(args.dsn, autocommit=True)
admin.execute("SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = 'cdc_lab'")
admin.execute("DROP TABLE IF EXISTS items")
admin.execute("CREATE TABLE items (id int PRIMARY KEY, status text, updated_at timestamptz NOT NULL DEFAULT clock_timestamp())")
admin.execute("INSERT INTO items SELECT g, 'new', clock_timestamp() FROM generate_series(1, 5) g")
time.sleep(0.2)
admin.execute("SELECT pg_create_logical_replication_slot('cdc_lab', 'test_decoding')")

seen, stop = [], threading.Event()
def poller():
    con = psycopg.connect(args.dsn, autocommit=True)
    last = con.execute("SELECT max(updated_at) FROM items").fetchone()[0]
    while not stop.is_set():
        for rid, status, ts in con.execute("SELECT id, status, updated_at FROM items WHERE updated_at > %s ORDER BY updated_at", (last,)):
            seen.append({"id": rid, "status": status}); last = max(last, ts)
        stop.wait(args.poll_every)
th = threading.Thread(target=poller); th.start()
w = psycopg.connect(args.dsn, autocommit=True)
time.sleep(1.5)
# (3) the slow transaction: now() is the transaction's start time, and it commits much later
slow = psycopg.connect(args.dsn)
slow.execute("UPDATE items SET status = 'slow', updated_at = now() WHERE id = 3")
# (1) three updates between two polls
for s in ("paid", "packed", "shipped"):
    w.execute("UPDATE items SET status = %s, updated_at = clock_timestamp() WHERE id = 1", (s,)); time.sleep(0.05)
time.sleep(args.poll_every * 2)
# (2) a delete
w.execute("DELETE FROM items WHERE id = 2")
time.sleep(args.poll_every * 2)
w.execute("UPDATE items SET status = 'later', updated_at = clock_timestamp() WHERE id = 4")   # moves the high-water mark on
time.sleep(max(0, args.slow_commit - 4 * args.poll_every))
slow.commit()                                         # commits now, with a timestamp from 10 s ago
time.sleep(args.poll_every * 3)
stop.set(); th.join()

changes = [r[0] for r in admin.execute("SELECT data FROM pg_logical_slot_get_changes('cdc_lab', NULL, NULL)")]
log = [c for c in changes if c.startswith("table public.items")]
actual = [("row 1, three updates", "statuses " + ", ".join(s["status"] for s in seen if s["id"] == 1) or "none",
           sum("id[integer]:1 " in c and "UPDATE" in c for c in log)),
          ("row 2, deleted", "seen" if any(s["id"] == 2 for s in seen) else "not seen", sum("DELETE" in c for c in log)),
          ("row 3, slow commit", "seen" if any(s["id"] == 3 for s in seen) else "missed", sum("id[integer]:3 " in c for c in log))]
rows = [{"action": a, "poller": pv, "log_events": n} for a, pv, n in actual]
for r in rows: print(f"{r['action']:<22} poller: {r['poller']:<34} logical decoding: {r['log_events']} event(s)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"actions": rows}}, f)