# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Idempotent sink with a conditional insert.

Records written to SQLite by a consumer that crashes and redelivers its last
batch: a plain INSERT into a table without a key, a PRIMARY KEY with
INSERT ... ON CONFLICT DO NOTHING, and dedup done in application code as
"SELECT, then INSERT if absent" by two consumers at once.
"""
import argparse, json, os, random, sqlite3, tempfile, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=10_000, help="Records.")
p.add_argument("--batch", type=int, default=100, help="Records per commit (the duplicate blast radius on a crash).")
p.add_argument("--crashes", type=int, default=10, help="Crashes during the run.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)


def consume(db, insert_sql):
    """Apply records in batches; each crash happens after a batch is written but before its offset is committed,
    so the batch is delivered again on restart."""
    crash_at = set(random.sample(range(args.batch, args.records, args.batch), args.crashes))
    pos = 0
    while pos < args.records:
        batch = range(pos, min(pos + args.batch, args.records))
        db.executemany(insert_sql, [(r, f"v{r}") for r in batch]); db.commit()
        if pos in crash_at:
            crash_at.discard(pos); continue                      # crash: offset not committed, redeliver
        pos += args.batch


def check_then_insert(path, ids, barrier):
    db = sqlite3.connect(path, timeout=30, isolation_level=None)
    barrier.wait()
    for r in ids:
        if db.execute("SELECT 1 FROM t WHERE id = ?", (r,)).fetchone() is None:   # check ...
            time.sleep(0.0002)
            db.execute("INSERT INTO t VALUES (?, ?)", (r, f"v{r}"))              # ... then act
    db.close()


rows = []
for name, ddl, sql in (("plain INSERT, no key", "CREATE TABLE t (id INTEGER, v TEXT)", "INSERT INTO t VALUES (?, ?)"),
                       ("PRIMARY KEY + ON CONFLICT DO NOTHING", "CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)",
                        "INSERT INTO t VALUES (?, ?) ON CONFLICT DO NOTHING")):
    db = sqlite3.connect(":memory:"); db.execute(ddl)
    consume(db, sql)
    total, distinct = db.execute("SELECT COUNT(*), COUNT(DISTINCT id) FROM t").fetchone()
    rows.append({"sink": name, "rows": total, "duplicates": total - distinct})

path = os.path.join(tempfile.mkdtemp(), "race.db")
db = sqlite3.connect(path); db.execute("CREATE TABLE t (id INTEGER, v TEXT)"); db.execute("PRAGMA journal_mode=WAL"); db.commit(); db.close()
barrier = threading.Barrier(2)
ids = list(range(2000))
ps = [threading.Thread(target=check_then_insert, args=(path, ids, barrier)) for _ in range(2)]   # two consumers, own connections
[x.start() for x in ps]; [x.join() for x in ps]
db = sqlite3.connect(path)
total, distinct = db.execute("SELECT COUNT(*), COUNT(DISTINCT id) FROM t").fetchone()
rows.append({"sink": "SELECT then INSERT, two consumers", "rows": total, "duplicates": total - distinct})
for r in rows: print(f"{r['sink']:<38} {r['rows']:>6} rows, {r['duplicates']:>5} duplicates")

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