# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Commit ordering, and the offset in the same transaction.

A consumer adds the numbers 1..10,000 to a total in SQLite and keeps a
position; it crashes (os._exit) after a random number of steps (any write, so
also between saving the position and processing the record) and
restarts, 20 times. (a) position saved to a file before processing, (b) after
processing, (c) total and position in one transaction. The right answer is
50,005,000. Then (b) with different commit intervals.
"""
import argparse, json, os, random, sqlite3, subprocess, sys, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=10_000, help="Numbers 1..N.")
p.add_argument("--crashes", type=int, default=20, help="Crashes per variant.")
p.add_argument("--intervals", type=lambda s: [int(v) for v in s.split(",")], default=[1, 10, 100], help="Commit intervals for variant (b).")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--worker", help=argparse.SUPPRESS)
p.add_argument("--dir", help=argparse.SUPPRESS)
p.add_argument("--crash-after", type=int, help=argparse.SUPPRESS)
p.add_argument("--every", type=int, default=1, help=argparse.SUPPRESS)
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

if args.worker:
    db = sqlite3.connect(os.path.join(args.dir, "db.sqlite"), isolation_level=None)
    posf = os.path.join(args.dir, "pos.txt")
    if args.worker == "c": pos = db.execute("SELECT pos FROM progress").fetchone()[0]
    else: pos = int(open(posf).read()) if os.path.exists(posf) else 0
    steps = [0]
    def step():                                                        # the process can die after any single step
        steps[0] += 1
        if args.crash_after and steps[0] >= args.crash_after: os._exit(1)
    for n in range(pos + 1, args.records + 1):
        if args.worker == "a":
            open(posf, "w").write(str(n)); step(); db.execute("UPDATE total SET v = v + ?", (n,)); step()
        elif args.worker == "b":
            db.execute("UPDATE total SET v = v + ?", (n,)); step()
            if n % args.every == 0 or n == args.records: open(posf, "w").write(str(n)); step()
        else:
            db.execute("BEGIN"); db.execute("UPDATE total SET v = v + ?", (n,)); step()
            db.execute("UPDATE progress SET pos = ?", (n,)); step(); db.execute("COMMIT"); step()
    sys.exit(0)

def variant(v, every=1):
    random.seed(args.seed)
    d = tempfile.mkdtemp(); db = sqlite3.connect(os.path.join(d, "db.sqlite"))
    db.execute("CREATE TABLE total (v INTEGER)"); db.execute("INSERT INTO total VALUES (0)")
    db.execute("CREATE TABLE progress (pos INTEGER)"); db.execute("INSERT INTO progress VALUES (0)"); db.commit(); db.close()
    for _ in range(args.crashes):
        subprocess.run([sys.executable, __file__, "--worker", v, "--dir", d, "--records", str(args.records), "--every", str(every),
                        "--crash-after", str(random.randint(1, 2 * args.records // args.crashes))])
    subprocess.run([sys.executable, __file__, "--worker", v, "--dir", d, "--records", str(args.records), "--every", str(every)], check=True)
    return sqlite3.connect(os.path.join(d, "db.sqlite")).execute("SELECT v FROM total").fetchone()[0]

right = args.records * (args.records + 1) // 2
rows = []
for v, label in (("a", "(a) position before processing"), ("b", "(b) position after processing"), ("c", "(c) both in one transaction")):
    t = variant(v); rows.append({"variant": label, "total": t, "off_by": t - right})
    print(f"{label:<34} total {t:>12,}   off by {t - right:+,}")
intervals = []
for every in args.intervals:
    t = variant("b", every); intervals.append({"commit_every": every, "off_by": t - right})
    print(f"(b) with the position saved every {every:>3} records: off by {t - right:+,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"variants": rows, "intervals": intervals, "summary": [{"right": right}]}}, f)