# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]", "pymysql"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Same name, different promise: compare engines.

One phantom test at the level both PostgreSQL and MySQL/InnoDB call
"Repeatable Read": T1 counts rows with x > 10, T2 inserts x = 11 and
commits, T1 counts again (plain and locking read) and inserts based on its
first count. Then: does a locking read by T1 block T2's insert?
Needs both services from the compose file (ports 55463 and 55464).
"""
import argparse, json, threading, time
import psycopg, pymysql

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--pg", default="postgresql://postgres:lab@localhost:55463/postgres", help="PostgreSQL connection string.")
p.add_argument("--mysql-port", type=int, default=55464, help="MySQL port.")
p.add_argument("--block-timeout", type=float, default=3.0, help="Seconds to wait before calling T2's insert blocked.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()


class Engine:
    def __init__(self, name):
        self.name = name
    def conn(self):
        if self.name == "PostgreSQL":
            c = psycopg.connect(args.pg); return c
        return pymysql.connect(host="127.0.0.1", port=args.mysql_port, user="root", password="lab", database="lab",
                               autocommit=False)
    def q(self, c, sql, *a):
        cur = c.cursor(); cur.execute(sql, a or None); return cur
    def begin_rr(self, c):
        if self.name == "PostgreSQL": self.q(c, "BEGIN ISOLATION LEVEL REPEATABLE READ")
        else: self.q(c, "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"); self.q(c, "START TRANSACTION")
    def setup(self):
        c = self.conn()
        self.q(c, "DROP TABLE IF EXISTS t")
        self.q(c, "CREATE TABLE t (id int PRIMARY KEY, x int)" if self.name == "PostgreSQL" else
               "CREATE TABLE t (id int PRIMARY KEY, x int, KEY (x)) ENGINE=InnoDB")
        for i, x in enumerate((5, 12, 15, 20), start=1): self.q(c, "INSERT INTO t VALUES (%s, %s)", i, x)
        c.commit(); c.close()


def phantom(e):
    e.setup()
    t1, t2 = e.conn(), e.conn()
    e.begin_rr(t1)
    first = e.q(t1, "SELECT count(*) FROM t WHERE x > 10").fetchone()[0]
    e.q(t2, "INSERT INTO t VALUES (100, 11)"); t2.commit()                       # T2 inserts x = 11, commits
    plain = e.q(t1, "SELECT count(*) FROM t WHERE x > 10").fetchone()[0]
    locking = len(e.q(t1, "SELECT id FROM t WHERE x > 10 FOR UPDATE").fetchall())
    e.q(t1, "INSERT INTO t VALUES (200, %s)", first)                             # a write based on the first count
    t1.commit()
    t1.close(); t2.close()
    return {"engine": e.name, "first_count": first, "second_count_plain": plain, "count_for_update": locking}


def blocking(e):
    e.setup()
    t1, t2 = e.conn(), e.conn()
    e.begin_rr(t1)
    e.q(t1, "SELECT id FROM t WHERE x > 10 FOR UPDATE").fetchall()               # a locking read of the range
    done = {}
    def insert():
        t = time.perf_counter()
        try:
            e.q(t2, "INSERT INTO t VALUES (300, 13)"); t2.commit()
        except Exception as ex:
            done["err"] = type(ex).__name__
        done["s"] = time.perf_counter() - t
    th = threading.Thread(target=insert); th.start(); th.join(args.block_timeout)
    blocked = th.is_alive()
    t1.commit(); th.join()
    t1.close(); t2.close()
    return {"blocked": "yes" if blocked else "no", "insert_waited_s": done["s"]}


rows = []
for e in (Engine("PostgreSQL"), Engine("MySQL InnoDB")):
    r = phantom(e) | blocking(e)
    rows.append(r)
    print(f"{e.name:<13} first count {r['first_count']}, second plain count {r['second_count_plain']}, "
          f"FOR UPDATE sees {r['count_for_update']}; after T1's locking read, T2's insert blocked: {r['blocked']} "
          f"(waited {r['insert_waited_s']:.2f} s)")

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