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

A 10-million-row table where country is 'US' for 95% of rows and 'LU' for
0.002%, indexed and analyzed. A prepared statement is executed with 'LU' six
or more times, then with 'US', under plan_cache_mode auto, force_custom_plan
and force_generic_plan: which plan each execution used, and its time.
Needs the PostgreSQL service from the compose file (port 55708).
"""
import argparse, json, re
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55708/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=10_000_000, help="Rows.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

c = psycopg.connect(args.dsn, autocommit=True, prepare_threshold=None)   # the driver must not prepare statements of its own
c.execute("DROP TABLE IF EXISTS t")
c.execute(f"""CREATE TABLE t AS SELECT g AS id, CASE WHEN g % 50000 = 0 THEN 'LU' WHEN g % 20 = 0 THEN 'FR' ELSE 'US' END AS country
              FROM generate_series(1, {args.rows}) g""")
c.execute("CREATE INDEX ON t (country)"); c.execute("VACUUM ANALYZE t")
rows = []
for mode in ("auto", "force_custom_plan", "force_generic_plan"):
    c.execute("DEALLOCATE ALL"); c.execute(f"SET plan_cache_mode = {mode}")
    c.execute("PREPARE q(text) AS SELECT count(*) FROM t WHERE country = $1")
    for i, val in enumerate(["LU"] * 7 + ["US", "LU"]):
        txt = "\n".join(r[0] for r in c.execute(f"EXPLAIN (ANALYZE) EXECUTE q('{val}')").fetchall())
        generic = "$1" in txt
        scan = "index-only" if "Index Only" in txt else "index" if "Index Scan" in txt or "Bitmap" in txt else "seq" if "Seq Scan" in txt else "?"
        rows.append({"mode": mode, "execution": i + 1, "value": val, "plan": ("generic" if generic else "custom") + f", {scan} scan",
                     "ms": float(re.search(r"Execution Time: ([\d.]+)", txt).group(1))})
for r in rows: print(f"{r['mode']:<18} #{r['execution']} {r['value']}: {r['plan']:<24} {r['ms']:8.2f} ms")

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