# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Watch the log retention problem happen.

A logical replication slot that nobody reads, while a writer inserts and
updates rows: the retained log grows until the slot is dropped.
Needs the PostgreSQL service from the compose file (port 55432).
"""
import argparse, json, threading, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55432/postgres", help="Connection string.")
p.add_argument("--minutes", type=float, default=4.0, help="How long the writer runs with the slot unread.")
p.add_argument("--sample-every", type=float, default=15.0, help="Seconds between measurements of the retained log.")
p.add_argument("--disk-gb", type=float, default=100.0, help="Free disk to project the time-to-full against.")
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 = 'lab'")
admin.execute("DROP TABLE IF EXISTS orders")
admin.execute("CREATE TABLE orders (order_id bigserial PRIMARY KEY, status text, total numeric, updated_at timestamptz DEFAULT now())")
admin.execute("SELECT pg_create_logical_replication_slot('lab', 'test_decoding')")     # created, never read

RETAINED = """SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), restart_lsn) FROM pg_replication_slots WHERE slot_name = 'lab'"""
WAL_DIR = "SELECT sum(size) FROM pg_ls_waldir()"

stop = threading.Event()
def writer():
    con = psycopg.connect(args.dsn, autocommit=True)
    while not stop.is_set():
        with con.transaction():
            con.execute("INSERT INTO orders (status, total) SELECT 'placed', random() * 100 FROM generate_series(1, 200)")
            con.execute("UPDATE orders SET status = 'paid', updated_at = now() WHERE order_id % 50 = (extract(epoch FROM now())::int % 50)")

w = threading.Thread(target=writer); w.start()
samples, t0 = [], time.time()
while time.time() - t0 < args.minutes * 60:
    time.sleep(args.sample_every)
    retained = admin.execute(RETAINED).fetchone()[0]
    samples.append({"minute": (time.time() - t0) / 60, "retained_mb": float(retained) / 1e6,
                    "wal_dir_mb": float(admin.execute(WAL_DIR).fetchone()[0]) / 1e6})
    print(f"{samples[-1]['minute']:5.2f} min   retained by the slot: {samples[-1]['retained_mb']:8.1f} MB   "
          f"pg_wal on disk: {samples[-1]['wal_dir_mb']:8.1f} MB")
stop.set(); w.join()

growth = (samples[-1]["retained_mb"] - samples[0]["retained_mb"]) / (samples[-1]["minute"] - samples[0]["minute"])  # MB/min
before = samples[-1]["wal_dir_mb"]
admin.execute("SELECT pg_drop_replication_slot('lab')")
admin.execute("CHECKPOINT"); admin.execute("SELECT pg_switch_wal()"); admin.execute("CHECKPOINT")
after = float(admin.execute(WAL_DIR).fetchone()[0]) / 1e6
hours = args.disk_gb * 1e3 / growth / 60
print(f"\ngrowth while unread: {growth:,.1f} MB per minute = {growth * 60 / 1e3:,.1f} GB per hour")
print(f"slot dropped + checkpoint: pg_wal {before:,.1f} MB -> {after:,.1f} MB")
print(f"at this rate a {args.disk_gb:g} GB disk fills in {hours:,.1f} hours ({hours / 24:,.1f} days)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"samples": samples, "summary": [{
            "growth_mb_per_min": growth, "wal_before_mb": before, "wal_after_mb": after,
            "hours_to_full": hours, "days_to_full": hours / 24, "disk_gb": args.disk_gb}]}}, f)