# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Reproduce a fan trap and a chasm trap.

A join to a finer-grained table repeats each row of the coarser one, so a sum
over the join counts it once per match (fan trap). Two independent
one-to-many branches from the same entity multiply (chasm trap), and an inner
join drops entities that lack one branch. Aggregating each branch first fixes both.
"""
import argparse, json, sqlite3

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

db = sqlite3.connect(":memory:")
db.executescript("""
CREATE TABLE orders(order_id TEXT, freight INT);
CREATE TABLE lines(order_id TEXT, line_no INT);
INSERT INTO orders VALUES('A',10),('B',10),('C',10);
INSERT INTO lines VALUES('A',1),('B',1),('B',2),('C',1),('C',2),('C',3),('C',4);
CREATE TABLE ord(customer TEXT, amount INT);
CREATE TABLE tix(customer TEXT, ticket INT);
INSERT INTO ord VALUES('C1',100),('C1',200),('C1',300),('C3',50);
INSERT INTO tix VALUES('C1',1),('C1',2);
""")
q = lambda sql: db.execute(sql).fetchall()
rows = [
    {"case": "freight, summed over the join to lines", "result": q("SELECT SUM(o.freight) FROM orders o JOIN lines l USING(order_id)")[0][0], "correct": 30},
    {"case": "freight, summed on orders alone", "result": q("SELECT SUM(freight) FROM orders")[0][0], "correct": 30},
    {"case": "C1 orders, summed over the join to tickets", "result": q("SELECT SUM(amount) FROM ord JOIN tix USING(customer) WHERE customer='C1'")[0][0], "correct": 600},
    {"case": "customers in the orders x tickets inner join", "result": len(q("SELECT DISTINCT customer FROM ord JOIN tix USING(customer)")), "correct": 2},
]
fixed = q("""
WITH o AS (SELECT customer, SUM(amount) AS amount FROM ord GROUP BY customer),
     t AS (SELECT customer, COUNT(*) AS tickets FROM tix GROUP BY customer)
SELECT o.customer, o.amount, COALESCE(t.tickets, 0) FROM o LEFT JOIN t USING(customer) ORDER BY 1""")
rows.append({"case": "C1 orders, one aggregate per branch", "result": fixed[0][1], "correct": 600})
rows.append({"case": "customers, one aggregate per branch", "result": len(fixed), "correct": 2})
for r in rows: print(f"{r['case']:<48} {r['result']:>6}   (correct: {r['correct']})")
print("per-branch result:", fixed)

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