File drop with manifest as commit

# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""File drop with the manifest as the commit.

A producer publishes 5 files the safe way (stage, checksum, atomic rename,
manifest last); a consumer processes a batch only if the manifest exists and
every checksum and the control total match. Then the batch is broken three
ways, and finally deleted along with its schedule.
"""
import argparse, hashlib, json, os, shutil, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=5, help="Files per batch.")
p.add_argument("--rows", type=int, default=1000, help="Rows per file.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def publish(root, batch, crash_after_renames=None, skip=None):
    stage, out = os.path.join(root, ".staging", batch), os.path.join(root, batch)
    os.makedirs(stage, exist_ok=True); os.makedirs(out, exist_ok=True)
    entries, total = [], 0
    for i in range(args.files):
        body = "".join(f"{batch},{i},{r},{r % 97}\n" for r in range(args.rows)).encode()
        name = f"part-{i}.csv"; open(os.path.join(stage, name), "wb").write(body)
        entries.append({"file": name, "sha256": hashlib.sha256(body).hexdigest(), "rows": args.rows}); total += args.rows
    for n, e in enumerate(entries):
        if crash_after_renames is not None and n == crash_after_renames: return                   # the producer dies mid-publish
        if skip == e["file"]: continue
        os.replace(os.path.join(stage, e["file"]), os.path.join(out, e["file"]))                   # atomic rename
    json.dump({"files": entries, "total_rows": total}, open(os.path.join(out, "_manifest.json.tmp"), "w"))
    os.replace(os.path.join(out, "_manifest.json.tmp"), os.path.join(out, "_manifest.json"))     # the commit

def consume(root, batch):
    out = os.path.join(root, batch); man = os.path.join(out, "_manifest.json")
    if not os.path.exists(man): return "rejected: no manifest (not committed)"
    m = json.load(open(man)); rows = 0
    for e in m["files"]:
        path = os.path.join(out, e["file"])
        if not os.path.exists(path): return f"rejected: {e['file']} missing"
        body = open(path, "rb").read()
        if hashlib.sha256(body).hexdigest() != e["sha256"]: return f"rejected: checksum mismatch in {e['file']}"
        rows += body.count(b"\n")
    return "accepted" if rows == m["total_rows"] else "rejected: control total mismatch"

root = tempfile.mkdtemp(); rows = []
publish(root, "b1"); rows.append({"case": "clean publish", "consumer": consume(root, "b1")})
publish(root, "b2", crash_after_renames=3); rows.append({"case": "producer killed between renames", "consumer": consume(root, "b2")})
publish(root, "b3"); path = os.path.join(root, "b3", "part-2.csv"); b = bytearray(open(path, "rb").read()); b[10] ^= 1; open(path, "wb").write(b)
rows.append({"case": "one byte flipped", "consumer": consume(root, "b3")})
publish(root, "b4", skip="part-4.csv"); rows.append({"case": "a file skipped", "consumer": consume(root, "b4")})
publish(root, "b5"); shutil.rmtree(os.path.join(root, "b5"))
seen = sorted(d for d in os.listdir(root) if not d.startswith("."))
rows.append({"case": "batch b5 deleted, no schedule", "consumer": f"consumer sees batches {', '.join(seen)}: nothing tells it b5 should exist"})
for r in rows: print(f"{r['case']:<34} {r['consumer']}")

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