# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The idempotency key and the unknown outcome.

SQLite wallets and a topups table keyed by an idempotency key. A top-up of 50
whose acknowledgment is lost is retried: once by a naive function (no key)
and once by a keyed one (insert the key, then update the balance, in one
transaction). Then the keys older than the retention window are deleted and
an old request is replayed.
"""
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()

def fresh():
    c = sqlite3.connect(":memory:", isolation_level=None)
    c.execute("CREATE TABLE wallets (owner TEXT PRIMARY KEY, balance INTEGER NOT NULL)")
    c.execute("CREATE TABLE topups (idem_key TEXT PRIMARY KEY, owner TEXT, amount INTEGER, at INTEGER)")
    c.execute("INSERT INTO wallets VALUES ('anna', 0)"); return c
def naive(c, owner, amount): c.execute("UPDATE wallets SET balance = balance + ? WHERE owner = ?", (amount, owner)); return "applied"
def topup(c, owner, amount, key, at=0):
    try:
        c.execute("BEGIN")
        c.execute("INSERT INTO topups VALUES (?, ?, ?, ?)", (key, owner, amount, at))           # fails on a duplicate key
        c.execute("UPDATE wallets SET balance = balance + ? WHERE owner = ?", (amount, owner)); c.execute("COMMIT"); return "applied"
    except sqlite3.IntegrityError:
        c.execute("ROLLBACK"); return "already applied"
bal = lambda c: c.execute("SELECT balance FROM wallets").fetchone()[0]
rows = []
c = fresh(); naive(c, "anna", 50); r2 = naive(c, "anna", 50)                                  # the first answer was lost; retry
rows.append({"version": "naive, no key", "second_call": r2, "balance": bal(c)})
c = fresh(); topup(c, "anna", 50, "req-1"); r2 = topup(c, "anna", 50, "req-1")
rows.append({"version": "idempotency key", "second_call": r2, "balance": bal(c)})
c.execute("DELETE FROM topups WHERE at < 7")                                                    # retention: keys older than a week are gone
r3 = topup(c, "anna", 50, "req-1", at=30)                                                       # an old request replayed a month later
rows.append({"version": "idempotency key, replay after the key expired", "second_call": r3, "balance": bal(c)})
for r in rows: print(f"{r['version']:<48} retry: {r['second_call']:<16} balance {r['balance']}")

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