# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Join elimination, with and without a constraint.

PostgreSQL: orders and stores with a primary key on stores.id and a foreign
key from orders.store_id. A LEFT JOIN that selects nothing from stores is
removed from the plan; an INNER JOIN is not. Then the trap of trusting a
constraint that is not enforced, shown by hand: an order whose store does
not exist makes the inner join and the join-free query give different counts.
Needs the PostgreSQL service from the compose file (port 55706).
"""
import argparse, json
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55706/postgres", help="Connection string.")
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 orders, stores")
c.execute("CREATE TABLE stores (id int PRIMARY KEY, name text)")
c.execute("CREATE TABLE orders (order_id int, store_id int REFERENCES stores (id))")
c.execute("INSERT INTO stores SELECT g, 's' || g FROM generate_series(1, 100) g")
c.execute("INSERT INTO orders SELECT g, 1 + g % 100 FROM generate_series(1, 100000) g"); c.execute("ANALYZE orders, stores")
def plan(sql): return [r[0] for r in c.execute("EXPLAIN (COSTS OFF) " + sql).fetchall()]
rows = []
for label, sql in (("LEFT JOIN, nothing selected from stores", "SELECT o.order_id FROM orders o LEFT JOIN stores s ON s.id = o.store_id"),
                   ("INNER JOIN, nothing selected from stores", "SELECT o.order_id FROM orders o JOIN stores s ON s.id = o.store_id")):
    pl = plan(sql); rows.append({"query": label, "stores_in_plan": "yes" if any("stores" in l for l in pl) else "no (join removed)",
                                 "plan": " / ".join(l.strip() for l in pl)})
# a constraint the engine does not enforce: drop the FK, insert an orphan, compare the inner join with the join-free answer
c.execute("ALTER TABLE orders DROP CONSTRAINT orders_store_id_fkey"); c.execute("INSERT INTO orders VALUES (100001, 999)")
with_join = c.execute("SELECT count(*) FROM orders o JOIN stores s ON s.id = o.store_id").fetchone()[0]
no_join = c.execute("SELECT count(*) FROM orders").fetchone()[0]
summary = {"with_join": with_join, "without_join": no_join, "differs": "yes" if with_join != no_join else "no"}
for r in rows: print(f"{r['query']:<42} stores in the plan: {r['stores_in_plan']}   [{r['plan']}]")
print(f"with an orphan order and an unenforced key: inner join counts {with_join:,}, the join removed counts {no_join:,}")

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