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

PostgreSQL with shared_buffers = 128 MB, a small hot table and a big table
about eight times the buffer pool. Warm the hot table with point lookups and
count its pages in the pool (pg_buffercache); read the whole big table with a
sequential scan and count again; then warm again and read the big table
through an index in random order (no ring buffer) and count again. Times a
pass of point lookups on the hot table after each read.
Needs the PostgreSQL service from the compose file (port 55601).
"""
import argparse, json, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55601/postgres", help="Connection string.")
p.add_argument("--big-rows", type=int, default=4_000_000, help="Rows in the big table (about 250 bytes each).")
p.add_argument("--hot-rows", type=int, default=20_000, help="Rows in the hot table.")
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("CREATE EXTENSION IF NOT EXISTS pg_buffercache")
c.execute("DROP TABLE IF EXISTS hot_table, big_table")
c.execute("CREATE TABLE hot_table (id int PRIMARY KEY, v text)")
c.execute(f"INSERT INTO hot_table SELECT g, repeat('h', 80) FROM generate_series(1, {args.hot_rows}) g")
c.execute("CREATE TABLE big_table (id int, rnd double precision, pad text)")
c.execute(f"INSERT INTO big_table SELECT g, random(), repeat('b', 200) FROM generate_series(1, {args.big_rows}) g")
c.execute("CREATE INDEX ON big_table (rnd)"); c.execute("VACUUM ANALYZE hot_table, big_table")

def q(sql): return c.execute(sql).fetchone()[0]
hot_pages = q("SELECT pg_relation_size('hot_table') / 8192")
big_mb = q("SELECT pg_relation_size('big_table')") / 2**20
pool_mb = q("SELECT setting::bigint * 8192 / 1048576 FROM pg_settings WHERE name = 'shared_buffers'")
def resident(): return q("""SELECT count(*) FROM pg_buffercache b JOIN pg_class c ON b.relfilenode = pg_relation_filenode(c.oid)
                            WHERE c.relname = 'hot_table' AND b.relforknumber = 0""")
def lookups():
    t = time.perf_counter()
    for i in range(1, args.hot_rows + 1, 7): c.execute("SELECT v FROM hot_table WHERE id = %s", (i,)).fetchone()
    return (time.perf_counter() - t) / len(range(1, args.hot_rows + 1, 7)) * 1e6
def warm():
    c.execute("SELECT count(v) FROM hot_table"); lookups(); lookups()

rows = []
warm(); rows.append({"step": "hot table warmed", "hot_pages_resident": resident(), "lookup_us": lookups()})
t = time.perf_counter(); q("SELECT count(*) FROM big_table"); seq_s = time.perf_counter() - t
rows.append({"step": "after a sequential scan of the big table", "hot_pages_resident": resident(), "lookup_us": lookups(), "read_s": seq_s})
warm(); rows.append({"step": "hot table warmed again", "hot_pages_resident": resident(), "lookup_us": lookups()})
c.execute("SET enable_seqscan = off"); c.execute("SET enable_bitmapscan = off")
plan = "\n".join(r[0] for r in c.execute("EXPLAIN SELECT count(pad) FROM big_table WHERE rnd >= 0").fetchall())
t = time.perf_counter(); q("SELECT count(pad) FROM big_table WHERE rnd >= 0"); idx_s = time.perf_counter() - t
rows.append({"step": "after reading the big table through an index (random order)", "hot_pages_resident": resident(), "lookup_us": lookups(), "read_s": idx_s})
for r in rows: r["hot_resident_pct"] = r["hot_pages_resident"] / hot_pages * 100
summary = {"hot_pages": hot_pages, "big_mb": big_mb, "pool_mb": pool_mb, "index_scan": "Index Scan" in plan}
print(f"buffer pool {pool_mb} MB, big table {big_mb:.0f} MB, hot table {hot_pages} pages")
for r in rows: print(f"{r['step']:<62} hot pages resident {r['hot_pages_resident']:>4}/{hot_pages} ({r['hot_resident_pct']:5.1f}%)   point lookup {r['lookup_us']:6.1f} µs")

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