# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Run an Elle-style check on your database.

A list-append workload: each transaction appends a unique integer to one
list and reads two others. Run at Read Committed, Repeatable Read and
Serializable with 8 clients, then rebuild the dependency graph from the
history (each list's final order is its write order) and look for cycles.
Needs the PostgreSQL service from the compose file (port 55462).
"""
import argparse, itertools, json, random, threading, time
from collections import defaultdict
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55462/postgres", help="Connection string.")
p.add_argument("--clients", type=int, default=8, help="Concurrent clients.")
p.add_argument("--seconds", type=float, default=30, help="Seconds of workload per isolation level.")
p.add_argument("--lists", type=int, default=4, help="Number of list rows.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

counter = itertools.count(1)


def workload(level):
    admin = psycopg.connect(args.dsn, autocommit=True)
    admin.execute("DROP TABLE IF EXISTS lists")
    admin.execute("CREATE TABLE lists (k int PRIMARY KEY, vals int[] NOT NULL DEFAULT '{}')")
    admin.execute(f"INSERT INTO lists SELECT k, '{{}}' FROM generate_series(1, {args.lists}) k")
    history, aborted, stop = [], [0], threading.Event()
    lock = threading.Lock()
    def client(cid):
        rnd = random.Random(args.seed * 100 + cid)
        con = psycopg.connect(args.dsn)
        while not stop.is_set():
            k = rnd.randint(1, args.lists); v = next(counter)
            reads = rnd.sample([x for x in range(1, args.lists + 1) if x != k], 2)
            try:
                con.execute(f"BEGIN ISOLATION LEVEL {level}")
                con.execute("UPDATE lists SET vals = array_append(vals, %s) WHERE k = %s", (v, k))
                got = {r: con.execute("SELECT vals FROM lists WHERE k = %s", (r,)).fetchone()[0] for r in reads}
                con.commit()
                with lock: history.append({"append": (k, v), "reads": got})
            except psycopg.errors.SerializationFailure:
                con.rollback()
                with lock: aborted[0] += 1
        con.close()
    ts = [threading.Thread(target=client, args=(c,)) for c in range(args.clients)]
    for t in ts: t.start()
    time.sleep(args.seconds); stop.set()
    for t in ts: t.join()
    final = {k: vals for k, vals in admin.execute("SELECT k, vals FROM lists").fetchall()}
    admin.close()
    return history, aborted[0], final


def check(history, final):
    """Edges from the list-append history (Elle's inference), then cycles classified by their rw edges."""
    writer = {h["append"][1]: i for i, h in enumerate(history)}
    succ = defaultdict(set); kind = {}
    def edge(a, b, k):
        if a != b and a is not None and b is not None:
            succ[a].add(b); kind.setdefault((a, b), set()).add(k)
    for k, order in final.items():
        for x, y in zip(order, order[1:]): edge(writer.get(x), writer.get(y), "ww")
    bad_reads = 0
    for i, h in enumerate(history):
        for k, seen in h["reads"].items():
            order = final[k]
            if seen != order[: len(seen)]: bad_reads += 1; continue          # not a prefix: an impossible read
            if seen: edge(writer.get(seen[-1]), i, "wr")
            if len(seen) < len(order): edge(i, writer.get(order[len(seen)]), "rw")
    # strongly connected components (iterative Tarjan)
    index, low, on, stack, sccs, n = {}, {}, set(), [], [], [0]
    for root in range(len(history)):
        if root in index: continue
        work = [(root, iter(succ[root]))]; index[root] = low[root] = n[0]; n[0] += 1; stack.append(root); on.add(root)
        while work:
            v, it = work[-1]
            w = next(it, None)
            if w is None:
                work.pop()
                if work: low[work[-1][0]] = min(low[work[-1][0]], low[v])
                if low[v] == index[v]:
                    comp = []
                    while True:
                        x = stack.pop(); on.discard(x); comp.append(x)
                        if x == v: break
                    if len(comp) > 1: sccs.append(set(comp))
            elif w not in index:
                index[w] = low[w] = n[0]; n[0] += 1; stack.append(w); on.add(w); work.append((w, iter(succ[w])))
            elif w in on:
                low[v] = min(low[v], index[w])
    # classify: a cycle with no rw edge (G0/G1c), exactly one (G-single, read skew), or more (G2, write skew)
    counts = {"G0/G1c": 0, "G-single": 0, "G2": 0}
    for comp in sccs:
        dep = {a: {b for b in succ[a] if b in comp and kind[(a, b)] - {"rw"}} for a in comp}
        def reach(src, dst):
            seen, todo = {src}, [src]
            while todo:
                x = todo.pop()
                if x == dst: return True
                for y in dep[x]:
                    if y not in seen: seen.add(y); todo.append(y)
            return False
        if any(reach(b, a) for a in comp for b in dep[a]): counts["G0/G1c"] += 1
        elif any(reach(b, a) for a in comp for b in succ[a] if b in comp and "rw" in kind[(a, b)]): counts["G-single"] += 1
        else: counts["G2"] += 1
    return counts, len(sccs), bad_reads


rows = []
for level in ("READ COMMITTED", "REPEATABLE READ", "SERIALIZABLE"):
    history, aborted, final = workload(level)
    counts, n_scc, bad = check(history, final)
    rows.append({"level": level, "committed": len(history), "aborted": aborted, "cyclic_components": n_scc, **counts,
                 "non_prefix_reads": bad})
    print(f"{level:<16} {len(history):>6} committed, {aborted:>5} aborted   cycles: G-single {counts['G-single']:>4}, "
          f"G2 {counts['G2']:>4}, G0/G1c {counts['G0/G1c']}   impossible reads {bad}")

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