# /// script
# requires-python = ">=3.10"
# dependencies = ["flask", "requests"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Idempotency: a webhook that retries.

A server that charges a customer and answers slowly (3 s) on every third
request; a client with a 1 s timeout that retries up to 3 times. Charges
without and with an Idempotency-Key, and the keys the server must keep.
"""
import argparse, json, logging, threading, time, uuid
import requests
from flask import Flask, request
from werkzeug.serving import make_server

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--requests", type=int, default=60, help="Distinct charges.")
p.add_argument("--slow-every", type=int, default=3, help="Every n-th request is slow.")
p.add_argument("--port", type=int, default=8142, help="Port.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
logging.getLogger("werkzeug").setLevel(logging.ERROR)
app = Flask(__name__); st = {"charges": 0, "n": 0, "keys": {}}; lock = threading.Lock()
@app.post("/charge")
def charge():
    key = request.headers.get("Idempotency-Key")
    with lock:
        st["n"] += 1; slow = st["n"] % args.slow_every == 0
        if key and key in st["keys"]: return st["keys"][key]
        st["charges"] += 1
        if key: st["keys"][key] = "charged"
    if slow: time.sleep(3)
    return "charged"
srv = make_server("127.0.0.1", args.port, app, threaded=True); threading.Thread(target=srv.serve_forever, daemon=True).start()
url = f"http://127.0.0.1:{args.port}/charge"
rows = []
for use_key in (False, True):
    with lock: st["charges"], st["n"] = 0, 0; st["keys"].clear()
    s = requests.Session()
    for i in range(args.requests):
        key = str(uuid.uuid4())
        for attempt in range(4):
            try: s.post(url, headers={"Idempotency-Key": key} if use_key else {}, timeout=1); break
            except requests.exceptions.Timeout: pass
    time.sleep(3.5)
    with lock: rows.append({"client": "with Idempotency-Key" if use_key else "no key", "requests": args.requests, "charges": st["charges"], "keys_stored": len(st["keys"])})
    print(f"{rows[-1]['client']:<22} {args.requests} requests -> {rows[-1]['charges']} charges   keys stored {rows[-1]['keys_stored']}")
srv.shutdown()

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