# /// script
# requires-python = ">=3.10"
# dependencies = ["pymysql", "cryptography"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Phantoms and gap locks depend on the plan.

MySQL/InnoDB at its default REPEATABLE READ. Session 1 runs SELECT ... WHERE
price BETWEEN 100 AND 120 FOR UPDATE and holds it; session 2 tries to insert
a seat priced inside the range and one priced far outside it. With an index on
price, then without, records which inserts wait (lock wait timeout of 2 s) and
counts the locks in performance_schema.data_locks.
Needs the MySQL service from the compose file (port 55810).
"""
import argparse, json, time
import pymysql

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--port", type=int, default=55810, help="MySQL port.")
p.add_argument("--rows", type=int, default=10_000, help="Seats.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def conn(): return pymysql.connect(host="127.0.0.1", port=args.port, user="root", password="lab", database="lab", autocommit=False)
admin = conn(); cur = admin.cursor()
rows = []
for indexed in (True, False):
    cur.execute("DROP TABLE IF EXISTS seats")
    cur.execute("CREATE TABLE seats (id int PRIMARY KEY, section char(1), price int" + (", INDEX (price)" if indexed else "") + ") ENGINE=InnoDB")
    cur.executemany("INSERT INTO seats VALUES (%s, %s, %s)", [(i, "ABCDE"[i % 5], (i * 37) % 1000) for i in range(1, args.rows + 1)]); admin.commit()
    s1 = conn(); c1 = s1.cursor(); c1.execute("BEGIN"); c1.execute("SELECT * FROM seats WHERE price BETWEEN 100 AND 120 FOR UPDATE"); matched = len(c1.fetchall())
    cur.execute("SELECT lock_type, lock_mode, count(*) FROM performance_schema.data_locks GROUP BY lock_type, lock_mode"); locks = cur.fetchall(); admin.commit()
    n_record = sum(n for t, m, n in locks if t == "RECORD")
    for label, price, sid in (("insert at price 110 (inside the range)", 110, 99_001), ("insert at price 900 (far outside)", 900, 99_002)):
        s2 = conn(); c2 = s2.cursor(); c2.execute("SET innodb_lock_wait_timeout = 2"); t = time.perf_counter()
        try: c2.execute("INSERT INTO seats VALUES (%s, 'F', %s)", (sid, price)); outcome = "inserted"; s2.rollback()
        except pymysql.err.OperationalError as e: outcome = "waited, then lock wait timeout" if e.args[0] == 1205 else str(e)
        rows.append({"index_on_price": "yes" if indexed else "no", "insert": label, "outcome": outcome, "waited_s": time.perf_counter() - t,
                     "record_locks": n_record, "rows_matched": matched})
        s2.close()
    s1.rollback(); s1.close()
for r in rows: print(f"index on price: {r['index_on_price']:<3}  {r['insert']:<40} {r['outcome']:<32} ({r['waited_s']:.1f} s)  record locks held: {r['record_locks']:,}")

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