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

A PostgreSQL table partitioned by day, with statistics on the old partitions.
A new day's partition is loaded and queried (filter on its date, joined to a
dimension) before and after ANALYZE on that partition only: the row estimate
and the join method. Also reads n_mod_since_analyze and computes when
auto-analyze would fire (0.1 × n + 50).
Needs the PostgreSQL service from the compose file (port 55702).
"""
import argparse, json, re
import psycopg

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

c = psycopg.connect(args.dsn, autocommit=True)
c.execute("DROP TABLE IF EXISTS events, stores")
c.execute("CREATE TABLE stores (id int PRIMARY KEY, name text)"); c.execute("INSERT INTO stores SELECT g, 's' || g FROM generate_series(1, 1000) g")
c.execute("CREATE TABLE events (day date, store_id int, amount numeric) PARTITION BY RANGE (day)")
for d in range(1, 6):
    c.execute(f"CREATE TABLE events_{d} PARTITION OF events FOR VALUES FROM ('2025-03-0{d}') TO ('2025-03-0{d + 1}')")
for d in range(1, 5):
    c.execute(f"INSERT INTO events_{d} SELECT DATE '2025-03-0{d}', 1 + g % 1000, g % 97 FROM generate_series(1, {args.rows_per_day}) g")
c.execute("ANALYZE stores"); [c.execute(f"ANALYZE events_{d}") for d in range(1, 5)]
c.execute("ALTER TABLE events_5 SET (autovacuum_enabled = off)")                     # keep auto-analyze from racing the lab
c.execute(f"INSERT INTO events_5 SELECT DATE '2025-03-05', 1 + g % 1000, g % 97 FROM generate_series(1, {args.rows_per_day}) g")
Q = "SELECT s.name, sum(e.amount) FROM events e JOIN stores s ON s.id = e.store_id WHERE e.day = DATE '2025-03-05' GROUP BY s.name"
def plan():
    txt = "\n".join(r[0] for r in c.execute("EXPLAIN ANALYZE " + Q).fetchall())
    scan = re.search(r"Seq Scan on events_5 e\s+\(cost=[^)]*rows=(\d+)[^)]*\) \(actual time=[^)]*rows=(\d+)", txt)
    join = "hash join" if "Hash Join" in txt else "nested loop" if "Nested Loop" in txt else "merge join" if "Merge Join" in txt else "?"
    t = float(re.search(r"Execution Time: ([\d.]+)", txt).group(1))
    return int(scan.group(1)), int(scan.group(2)), join, t
c.execute("SELECT pg_stat_force_next_flush()"); c.execute("SELECT 1"); import time; time.sleep(1)
mod = c.execute("SELECT n_mod_since_analyze, n_live_tup FROM pg_stat_user_tables WHERE relname = 'events_5'").fetchone()
rows = []
e, a, j, t = plan(); rows.append({"state": "new partition, never analyzed", "estimate": e, "actual": a, "join": j, "ms": t})
c.execute("ANALYZE events_5")
e, a, j, t = plan(); rows.append({"state": "after ANALYZE events_5", "estimate": e, "actual": a, "join": j, "ms": t})
summary = {"n_mod_since_analyze": mod[0], "trigger_at": 0.1 * args.rows_per_day + 50, "note": "threshold = 0.1 × reltuples + 50; reltuples is unknown (-1) before the first analyze"}
for r in rows: print(f"{r['state']:<32} estimate {r['estimate']:>8,}  actual {r['actual']:>8,}  join: {r['join']:<12} {r['ms']:.1f} ms")
print(f"n_mod_since_analyze = {mod[0]:,}; auto-analyze fires after 0.1 × n + 50 = {summary['trigger_at']:,.0f} changes")

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