"""
03_deadlock_detect.py
---------------------
故意构造一个死锁，观察 PostgreSQL 自动检测并回滚一个事务：
  事务 A: 锁住 account 1 -> 想锁 account 2
  事务 B: 锁住 account 2 -> 想锁 account 1
  PG 在 deadlock_timeout(默认 1s) 后检测到环，回滚其中一个，抛 SQLSTATE 40P01。

依赖: pip install psycopg[binary]>=3.1
运行: python 03_deadlock_detect.py
"""
import threading
import time
import psycopg
from psycopg import errors

DSN = "host=127.0.0.1 port=5432 dbname=learn_pg user=postgres"


def worker(name: str, first_id: int, second_id: int, barrier: threading.Barrier,
           result: dict):
    try:
        with psycopg.connect(DSN) as conn:
            with conn.transaction():
                with conn.cursor() as cur:
                    cur.execute(
                        "UPDATE ch11_bank_account SET balance = balance - 1 WHERE id = %s",
                        (first_id,),
                    )
                    print(f"[{name}] 已锁 account {first_id}")
                    barrier.wait(timeout=5)  # 等对方也拿到第一把锁
                    print(f"[{name}] 准备锁 account {second_id} ...")
                    cur.execute(
                        "UPDATE ch11_bank_account SET balance = balance + 1 WHERE id = %s",
                        (second_id,),
                    )
                    print(f"[{name}] 拿到 account {second_id}, 提交")
        result[name] = "OK"
    except errors.DeadlockDetected as e:
        print(f"[{name}] 被 PG 选为 victim，回滚: {e.diag.message_primary}")
        result[name] = "DEADLOCK"
    except Exception as e:
        print(f"[{name}] 其他异常: {type(e).__name__}: {e}")
        result[name] = "ERROR"


def main():
    barrier = threading.Barrier(2)
    result: dict = {}
    a = threading.Thread(target=worker, args=("A", 1, 2, barrier, result))
    b = threading.Thread(target=worker, args=("B", 2, 1, barrier, result))
    t0 = time.time()
    a.start(); b.start()
    a.join();  b.join()
    print(f"\n--- 结束，耗时 {time.time()-t0:.2f}s, 结果: {result} ---")
    print("（如果 deadlock_timeout = 1s，预期约 ~1s 后看到回滚）")


if __name__ == "__main__":
    main()