#!/usr/bin/env python3
"""Loading a batch job's output into a serving store: row-by-row vs bulk, DDIA 2e Ch. 11.

A batch job has produced N derived records and now has to load them into a
serving store (here: a fresh on-disk sqlite database). Two ways to do it:

  A) Row-by-row: one INSERT per record, each in its OWN transaction
     (conn.execute(INSERT) then conn.commit(), every row). This is the naive
     "just call the DB client inside the job's inner loop" pattern. Every commit
     forces a durable write -- an fsync to stable storage -- so the job pays that
     durability tax N times.

  B) Bulk: all N rows in ONE transaction via conn.executemany(), a single
     commit(). One fsync amortized over every row.

The aha (magnitude): B is not "a bit" faster, it is orders of magnitude faster,
because A pays ~N fsyncs and B pays ~1. We measure the real ratio.

The secondary aha (causation): row-by-row writes mutate the LIVE serving table
incrementally, so a job that fails partway leaves partial output behind; the
retry re-inserts from the start and DUPLICATES the rows it already wrote. The
batch-job discipline -- write to a FRESH table, then atomically swap it in -- keeps
the all-or-nothing guarantee: a crash leaves the old table untouched and the
retry is clean.

Pure standard library (sqlite3). Deterministic data (seed 42). Timings are real
wall-clock and vary run-to-run; the large RATIO is the reproducible result.
"""
import os
import random
import sqlite3
import tempfile
import time

N = 50_000            # records the batch job produced
CRASH_AT = 30_000     # row-by-row job "crashes" after writing this many


def make_records(n):
    """Deterministic derived records: (user_id, score, bucket)."""
    rng = random.Random(42)
    return [(i, rng.randint(0, 1_000_000), rng.choice("ABCDEF")) for i in range(n)]


def fresh_db(path):
    """A brand-new database file with an empty target table."""
    if os.path.exists(path):
        os.remove(path)
    conn = sqlite3.connect(path)
    conn.execute("CREATE TABLE scores (user_id INTEGER PRIMARY KEY, score INTEGER, bucket TEXT)")
    conn.commit()
    return conn


# --- Method A: row-by-row, one durable commit per record (the slow path) ---

def load_row_by_row(path, rows):
    conn = fresh_db(path)
    start = time.perf_counter()
    for r in rows:
        conn.execute("INSERT INTO scores VALUES (?, ?, ?)", r)
        conn.commit()          # a durable commit -- fsync -- for EVERY row
    elapsed = time.perf_counter() - start
    count = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.close()
    return elapsed, count


# --- Method B: one transaction, executemany, a single durable commit ---

def load_bulk(path, rows):
    conn = fresh_db(path)
    start = time.perf_counter()
    conn.executemany("INSERT INTO scores VALUES (?, ?, ?)", rows)
    conn.commit()              # ONE fsync for all N rows
    elapsed = time.perf_counter() - start
    count = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.close()
    return elapsed, count


# --- Partial-output aha: a row-by-row job that crashes mid-load, then retries ---

def load_row_by_row_crashing(path, rows, crash_at):
    """Write into the LIVE table row-by-row; 'crash' after crash_at rows.

    Because each row was already committed, the partial output survives the crash.
    """
    conn = fresh_db(path)
    written = 0
    for r in rows:
        conn.execute("INSERT INTO scores VALUES (?, ?, ?)", r)
        conn.commit()
        written += 1
        if written >= crash_at:
            break              # simulate the process dying here
    survived = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.close()
    return survived


def retry_row_by_row_ignoring_partial(path, rows):
    """Retry re-inserts from the start into the SAME live table.

    INSERT OR IGNORE past the PRIMARY KEY would hide the double-write; a plain
    re-run of "the job" that used a non-unique key (e.g. an append-only events
    table) would duplicate. We model the honest failure: the rows already present
    collide on the PRIMARY KEY, so we count how many the retry could NOT re-insert
    -- i.e. how many are now stranded partial output that a clean run never wanted.
    """
    conn = sqlite3.connect(path)  # NB: same file, NOT fresh -- partial output is still here
    collisions = 0
    for r in rows:
        try:
            conn.execute("INSERT INTO scores VALUES (?, ?, ?)", r)
            conn.commit()
        except sqlite3.IntegrityError:
            collisions += 1       # this row was already written by the crashed run
    final = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.close()
    return collisions, final


def bulk_load_then_swap_crashing(path, rows, crash_at):
    """The disciplined path: build a FRESH staging table, swap it in atomically.

    We crash while building the staging table. Because the live table is never
    touched until the atomic swap, the crash leaves NO partial output in the
    serving table, and a retry rebuilds the staging table from scratch cleanly.
    """
    conn = fresh_db(path)  # live table 'scores' starts empty (imagine it holds yesterday's data)
    # Build staging in one transaction; 'crash' before we ever swap.
    conn.execute("CREATE TABLE scores_staging (user_id INTEGER PRIMARY KEY, score INTEGER, bucket TEXT)")
    partial = rows[:crash_at]
    conn.executemany("INSERT INTO scores_staging VALUES (?, ?, ?)", partial)
    # ---- process dies HERE, before commit + swap ----
    conn.close()

    # Retry from a fresh connection: the uncommitted staging work is gone, the
    # live 'scores' table is intact. Rebuild staging fully, then swap atomically.
    conn = sqlite3.connect(path)
    live_after_crash = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.execute("DROP TABLE IF EXISTS scores_staging")
    conn.execute("CREATE TABLE scores_staging (user_id INTEGER PRIMARY KEY, score INTEGER, bucket TEXT)")
    conn.executemany("INSERT INTO scores_staging VALUES (?, ?, ?)", rows)
    conn.commit()
    # Atomic swap: DROP + RENAME inside one transaction.
    conn.execute("DROP TABLE scores")
    conn.execute("ALTER TABLE scores_staging RENAME TO scores")
    conn.commit()
    final = conn.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
    conn.close()
    return live_after_crash, final


def main():
    rows = make_records(N)
    tmp = tempfile.mkdtemp(prefix="ddia_ch11_")
    path_a = os.path.join(tmp, "row_by_row.db")
    path_b = os.path.join(tmp, "bulk.db")
    path_c = os.path.join(tmp, "crash_rbr.db")
    path_d = os.path.join(tmp, "crash_swap.db")

    print(f"batch job produced {N:,} derived records; load them into a fresh sqlite serving store.\n")

    # --- Step 1: row-by-row, one durable commit per row ---
    t_rbr, count_rbr = load_row_by_row(path_a, rows)
    print("A) row-by-row: one INSERT per record, each in its OWN transaction (commit-per-row):")
    print(f"    loaded {count_rbr:,} rows in {t_rbr:.2f} s")
    print(f"    = {N / t_rbr:,.0f} rows/sec\n")

    # --- Step 2: bulk, one transaction ---
    t_bulk, count_bulk = load_bulk(path_b, rows)
    print("B) bulk: all rows in ONE transaction (executemany + single commit):")
    print(f"    loaded {count_bulk:,} rows in {t_bulk:.3f} s")
    print(f"    = {N / t_bulk:,.0f} rows/sec\n")

    # --- Step 3: the ratio ---
    ratio = t_rbr / t_bulk
    assert count_rbr == count_bulk == N, "both methods must load the same row count"
    print("ratio:")
    print(f"    row-by-row / bulk  =  {t_rbr:.2f} s / {t_bulk:.3f} s  ~=  {ratio:,.0f}x slower")
    print(f"    both loaded the identical {N:,} rows -- same result, ~{ratio:,.0f}x the time.\n")

    # --- Step 4: partial output on a crashed row-by-row load ---
    survived = load_row_by_row_crashing(path_c, rows, CRASH_AT)
    collisions, final_rbr = retry_row_by_row_ignoring_partial(path_c, rows)
    print("C) row-by-row load CRASHES after {:,} rows, then the job RETRIES:".format(CRASH_AT))
    print(f"    partial output surviving the crash (already-committed rows): {survived:,}")
    print(f"    on retry, rows that collide with that partial output:        {collisions:,}")
    print(f"    -> the crashed run leaked {collisions:,} rows the retry cannot cleanly re-write;")
    print(f"       row-by-row mutation of the live store is NOT idempotent under retry.\n")

    # --- Step 5: bulk-load-then-swap survives the same crash cleanly ---
    live_after_crash, final_swap = bulk_load_then_swap_crashing(path_d, rows, CRASH_AT)
    print("D) bulk-load-then-swap CRASHES at the same point, then RETRIES:")
    print(f"    partial output in the live serving table after the crash: {live_after_crash:,}")
    print(f"    after a clean retry (rebuild staging, atomic swap):        {final_swap:,} rows")
    print(f"    -> the crash left the live table untouched ({live_after_crash:,} rows); the retry is")
    print(f"       all-or-nothing and idempotent -- exactly {N:,} rows, no duplicates.")


if __name__ == "__main__":
    main()