# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "sqlglot"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Lineage coverage: static parse against runtime.

Twenty small jobs (written for this lab: plain SQL, branches that run only on
some days, and SQL built from strings) read and write tables in DuckDB.
Static lineage parses every SQL string literal in the jobs' source with
sqlglot. Runtime lineage runs the jobs and reads the tables each executed
statement actually scanned, from DuckDB's query profile. Compares the edges
each method finds.
"""
import argparse, ast, inspect, json, os, re, tempfile, textwrap
import duckdb, sqlglot
from sqlglot import exp

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--weekday", type=int, default=2, help="Day the jobs run on (0 = Monday); some branches run only on some days.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

SOURCES = ["orders", "customers", "payments", "refunds", "shipments", "products", "events_2025_03", "events_2025_04", "fx_rates", "regions"]

def j01(run): run("CREATE OR REPLACE TABLE daily_orders AS SELECT * FROM orders")
def j02(run): run("CREATE OR REPLACE TABLE order_customer AS SELECT o.*, c.region FROM orders o JOIN customers c USING (id)")
def j03(run): run("CREATE OR REPLACE TABLE net_payments AS SELECT * FROM payments EXCEPT SELECT * FROM refunds")
def j04(run): run("CREATE OR REPLACE TABLE late_shipments AS SELECT * FROM shipments WHERE x > 5")
def j05(run, day=None):
    if args.weekday == 0: run("CREATE OR REPLACE TABLE weekly_products AS SELECT * FROM products")    # Mondays only
    else: run("CREATE OR REPLACE TABLE weekly_products AS SELECT * FROM regions")
def j06(run):
    month = "2025_04"; run(f"CREATE OR REPLACE TABLE recent_events AS SELECT * FROM events_{month}")   # name built at run time
def j07(run): run("CREATE OR REPLACE TABLE fx_orders AS SELECT * FROM orders JOIN fx_rates USING (id)")
def j08(run):
    t = "pay" + "ments"; run("CREATE OR REPLACE TABLE pay_copy AS SELECT * FROM " + t)
def j09(run): run("CREATE OR REPLACE TABLE region_rollup AS SELECT region, count(*) n FROM customers GROUP BY region")
def j10(run):
    if False: run("CREATE OR REPLACE TABLE old_report AS SELECT * FROM refunds")                       # dead code
    run("CREATE OR REPLACE TABLE new_report AS SELECT * FROM shipments")
def j11(run): run("CREATE OR REPLACE TABLE cust_orders AS WITH x AS (SELECT * FROM orders) SELECT * FROM x JOIN customers USING (id)")
def j12(run):
    for m in ("2025_03", "2025_04"): run(f"INSERT INTO all_events SELECT * FROM events_{m}")
def j13(run): run("CREATE OR REPLACE TABLE prod_regions AS SELECT * FROM products, regions")
def j14(run):
    src = os.environ.get("J14_SOURCE", "orders"); run(f"CREATE OR REPLACE TABLE env_copy AS SELECT * FROM {src}")
def j15(run): run("CREATE OR REPLACE TABLE refund_view AS SELECT * FROM refunds WHERE id IN (SELECT id FROM payments)")
def j16(run):
    if args.weekday >= 5: run("CREATE OR REPLACE TABLE weekend_fx AS SELECT * FROM fx_rates")           # weekends only
def j17(run): run("CREATE OR REPLACE TABLE ship_customer AS SELECT * FROM shipments JOIN customers USING (id)")
def j18(run):
    cols = "id, x"; run("CREATE OR REPLACE TABLE slim_orders AS SELECT %s FROM orders" % cols)
def j19(run): run("CREATE OR REPLACE TABLE product_copy AS SELECT * FROM products")
def j20(run):
    tables = [t for t in SOURCES if t.startswith("events_")]
    run("CREATE OR REPLACE TABLE event_union AS " + " UNION ALL ".join(f"SELECT * FROM {t}" for t in tables))
JOBS = [j01, j02, j03, j04, j05, j06, j07, j08, j09, j10, j11, j12, j13, j14, j15, j16, j17, j18, j19, j20]

def tables_in(sql):
    try: tree = sqlglot.parse_one(sql, read="duckdb")
    except Exception: return None
    target = tree.this.find(exp.Table) if isinstance(tree, (exp.Create, exp.Insert)) else None
    ctes = {c.alias for c in tree.find_all(exp.CTE)}
    reads = {t.name for t in tree.find_all(exp.Table) if t is not target and t.name not in ctes}
    return (target.name if target else None), reads

static, unparsed = set(), 0
for job in JOBS:
    for node in ast.walk(ast.parse(textwrap.dedent(inspect.getsource(job)))):
        if isinstance(node, ast.Constant) and isinstance(node.value, str) and re.match(r"\s*(CREATE|INSERT|SELECT)", node.value):
            got = tables_in(node.value)
            if got is None or not got[0]: unparsed += 1; continue
            for src in got[1]: static.add((job.__name__, src, got[0]))

work = tempfile.mkdtemp(); c = duckdb.connect(os.path.join(work, "lab.duckdb"))
for t in SOURCES: c.execute(f"CREATE TABLE {t} AS SELECT i AS id, i % 10 AS x, 'r' || (i % 3) AS region FROM range(100) t(i)")
c.execute("CREATE TABLE all_events AS SELECT * FROM events_2025_03 LIMIT 0")
prof = os.path.join(work, "profile.json")
c.execute("PRAGMA enable_profiling = 'json'"); c.execute(f"PRAGMA profiling_output = '{prof}'")
runtime = set()
def scanned(node, out):
    info = node.get("extra_info") or {}
    t = info.get("Table") if isinstance(info, dict) else None
    if t: out.add(t.split(".")[-1])                                       # catalog.schema.table -> table
    for ch in node.get("children", []): scanned(ch, out)
for job in JOBS:
    def run(sql, job=job):
        c.execute(sql); target = tables_in(sql)[0]; reads = set(); scanned(json.load(open(prof)), reads)
        for src in reads - {target}: runtime.add((job.__name__, src, target))
    job(run)

both, only_static, only_runtime = static & runtime, static - runtime, runtime - static
rows = [{"method": "static parse", "edges": len(static), "coverage_pct": len(static) / len(static | runtime) * 100},
        {"method": "runtime log", "edges": len(runtime), "coverage_pct": len(runtime) / len(static | runtime) * 100}]
detail = [{"edge": f"{s} → {t}", "job": j, "found_by": "both" if (j, s, t) in both else "parse only" if (j, s, t) in only_static else "runtime only"}
          for j, s, t in sorted(static | runtime)]
summary = {"union": len(static | runtime), "both": len(both), "only_static": len(only_static), "only_runtime": len(only_runtime),
           "unparsed_strings": unparsed, "weekday": args.weekday}
print(f"{len(static | runtime)} edges in all: both {len(both)}, parse only {len(only_static)}, runtime only {len(only_runtime)} ({unparsed} SQL strings the parser could not resolve)")
for r in rows: print(f"  {r['method']:<13} {r['edges']:>3} edges, coverage {r['coverage_pct']:.0f}%")
for d in detail:
    if d["found_by"] != "both": print(f"  {d['found_by']:<13} {d['job']}: {d['edge']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"methods": rows, "edges": [d for d in detail if d['found_by'] != 'both']}}, f)