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

Single-row inserts into orders, and one bulk INSERT ... SELECT, with and
without the shadow-table trigger that records every change as JSON.
Needs the PostgreSQL service from the compose file (port 55433).
"""
import argparse, json, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55433/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=200_000, help="Rows inserted, one statement each, and in the bulk insert.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

SHADOW = """
CREATE TABLE orders_changes (
    change_id  bigserial PRIMARY KEY,
    op         char(1),            -- I, U or D
    order_id   bigint,
    old_row    jsonb,
    new_row    jsonb,
    changed_at timestamptz DEFAULT now()
);
CREATE FUNCTION capture_orders() RETURNS trigger AS $$
BEGIN
  IF TG_OP = 'INSERT' THEN
    INSERT INTO orders_changes(op, order_id, new_row)
    VALUES ('I', NEW.order_id, to_jsonb(NEW));
  ELSIF TG_OP = 'UPDATE' THEN
    INSERT INTO orders_changes(op, order_id, old_row, new_row)
    VALUES ('U', NEW.order_id, to_jsonb(OLD), to_jsonb(NEW));
  ELSE
    INSERT INTO orders_changes(op, order_id, old_row)
    VALUES ('D', OLD.order_id, to_jsonb(OLD));
  END IF;
  RETURN NULL;
END $$ LANGUAGE plpgsql;
CREATE TRIGGER orders_cdc AFTER INSERT OR UPDATE OR DELETE ON orders
FOR EACH ROW EXECUTE FUNCTION capture_orders();
"""

con = psycopg.connect(args.dsn, autocommit=True)

def reset(trigger):
    con.execute("DROP TABLE IF EXISTS orders, orders_changes CASCADE; DROP FUNCTION IF EXISTS capture_orders CASCADE")
    con.execute("CREATE TABLE orders (order_id bigint PRIMARY KEY, customer_id bigint, status text, total numeric)")
    if trigger: con.execute(SHADOW)

def single_rows():
    t = time.perf_counter()
    with con.cursor() as cur:
        for i in range(args.rows):                    # one statement, one commit per row (autocommit)
            cur.execute("INSERT INTO orders VALUES (%s, %s, 'placed', %s)", (i, i % 1000, i % 97))
    return time.perf_counter() - t

def bulk():
    t = time.perf_counter()
    con.execute(f"INSERT INTO orders SELECT i, i % 1000, 'placed', i % 97 FROM generate_series(1, {args.rows}) i")
    return time.perf_counter() - t

results = {}
for trigger in (False, True):
    reset(trigger); s = single_rows()
    reset(trigger); b = bulk()
    results[trigger] = (s, b)
    label = "with trigger" if trigger else "no trigger  "
    print(f"{label}: {args.rows:,} single-row inserts {s:6.1f} s ({s / args.rows * 1e6:5.0f} µs each)   "
          f"bulk INSERT ... SELECT {b:5.2f} s")
changes = con.execute("SELECT count(*) FROM orders_changes").fetchone()[0]
print(f"\nper statement the trigger costs {results[True][0] / results[False][0]:.1f}x; "
      f"on the bulk insert {results[True][1] / results[False][1]:.1f}x "
      f"({changes:,} shadow rows written, one per inserted row)")
ver = con.execute("SHOW server_version").fetchone()[0]

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "timings": [{"mode": m, "statement": st, "seconds": results[tr][i]}
                        for tr, m in ((False, "no trigger"), (True, "with trigger"))
                        for i, st in ((0, "single-row inserts"), (1, "bulk INSERT ... SELECT"))],
            "summary": [{"single_ratio": results[True][0] / results[False][0], "bulk_ratio": results[True][1] / results[False][1],
                         "single_us_plain": results[False][0] / args.rows * 1e6, "single_us_trigger": results[True][0] / args.rows * 1e6,
                         "shadow_rows": changes, "server_version": ver}]}}, f)