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

One task loads a day's 100 rows into SQLite and crashes after half of them,
then is retried. Four ways: append inside a transaction, append without one,
overwrite (delete the day's partition, then insert) in one transaction, and
overwrite with random amounts, to see an idempotent retry change the data.
"""
import argparse, json, random, sqlite3

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=100, help="Rows per day.")
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 case(mode, transactional, amount):
    db = sqlite3.connect(":memory:", isolation_level=None)
    db.execute("CREATE TABLE sales(day TEXT, store INT, amount REAL)")
    def load(day, crash=False):
        if transactional: db.execute("BEGIN")
        try:
            if mode == "overwrite": db.execute("DELETE FROM sales WHERE day=?", (day,))
            for s in range(args.rows):
                if crash and s == args.rows // 2: raise RuntimeError("boom")
                db.execute("INSERT INTO sales VALUES (?,?,?)", (day, s, amount()))
            if transactional: db.execute("COMMIT")
        except RuntimeError:
            if transactional: db.execute("ROLLBACK")
            raise
    try: load("2025-03-01", crash=True)
    except RuntimeError: pass
    after_crash = db.execute("SELECT count(*), sum(amount) FROM sales").fetchone()
    load("2025-03-01")                                                 # the retry
    first = db.execute("SELECT count(*), sum(amount) FROM sales").fetchone()
    load("2025-03-01")                                                 # and one more, as a rerun would
    second = db.execute("SELECT count(*), sum(amount) FROM sales").fetchone()
    return after_crash[0], first, second

rows = []
for label, mode, tx, amount in (("append, in a transaction", "append", True, lambda: 10.0),
                                ("append, no transaction", "append", False, lambda: 10.0),
                                ("overwrite, in a transaction", "overwrite", True, lambda: 10.0),
                                ("overwrite, random amounts", "overwrite", True, random.random)):
    crash_rows, first, second = case(mode, tx, amount)
    rows.append({"load": label, "rows_after_crash": crash_rows, "rows_after_retry": first[0], "rows_after_rerun": second[0],
                 "total_after_retry": round(first[1], 2), "total_after_rerun": round(second[1], 2),
                 "rerun_changed_total": "yes" if round(first[1], 6) != round(second[1], 6) else "no"})
for r in rows: print(f"{r['load']:<30} after crash {r['rows_after_crash']:>3} rows; after retry {r['rows_after_retry']:>3} (total {r['total_after_retry']}); after a rerun {r['rows_after_rerun']:>3} (total {r['total_after_rerun']}, changed: {r['rerun_changed_total']})")

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