"""
05_blocking_query.py
--------------------
故意制造一个阻塞场景，再用 `pg_blocking_pids()` + `pg_locks` 把
"谁在等谁、等多久、等什么 SQL" 全部打印出来 —— 模拟生产排查流程。

依赖: pip install psycopg[binary]>=3.1
运行: python 05_blocking_query.py
"""
import threading
import time
import psycopg

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

BLOCKING_CHAIN_SQL = """
SELECT
    blocked.pid                                         AS blocked_pid,
    blocked.usename                                     AS blocked_user,
    LEFT(blocked.query, 60)                             AS blocked_query,
    blocking.pid                                        AS blocking_pid,
    blocking.usename                                    AS blocking_user,
    LEFT(blocking.query, 60)                            AS blocking_query,
    blocked.wait_event_type                             AS wait_type,
    blocked.wait_event                                  AS wait_event,
    EXTRACT(EPOCH FROM (now() - blocked.xact_start))::INT AS blocked_secs
  FROM pg_stat_activity blocked
  JOIN pg_stat_activity blocking
    ON blocking.pid = ANY(pg_blocking_pids(blocked.pid))
 WHERE blocked.wait_event_type = 'Lock';
"""

LOCKS_SQL = """
SELECT pid,
       locktype,
       mode,
       granted,
       relation::regclass AS rel
  FROM pg_locks
 WHERE relation = 'ch11_seckill'::regclass
 ORDER BY granted DESC, pid;
"""


def long_holder(ready: threading.Event):
    with psycopg.connect(DSN) as conn:
        with conn.transaction():
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM ch11_seckill WHERE id = 1 FOR UPDATE")
                ready.set()
                print("[holder] 已锁住 ch11_seckill id=1，模拟业务跑 6s …")
                time.sleep(6)


def waiter(ready: threading.Event):
    ready.wait()
    print("[waiter] 我也想 UPDATE 同一行 …")
    with psycopg.connect(DSN, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute("UPDATE ch11_seckill SET stock = stock - 1 WHERE id = 1")
            print("[waiter] 终于拿到锁并更新成功")


def watcher(ready: threading.Event):
    ready.wait()
    time.sleep(1)
    with psycopg.connect(DSN, autocommit=True) as conn:
        with conn.cursor() as cur:
            print("\n--- 阻塞链 (pg_blocking_pids) ---")
            cur.execute(BLOCKING_CHAIN_SQL)
            for r in cur.fetchall():
                print(r)

            print("\n--- pg_locks on ch11_seckill ---")
            cur.execute(LOCKS_SQL)
            for r in cur.fetchall():
                print(r)


def main():
    ready = threading.Event()
    h = threading.Thread(target=long_holder, args=(ready,))
    w = threading.Thread(target=waiter, args=(ready,))
    s = threading.Thread(target=watcher, args=(ready,))
    for t in (h, w, s):
        t.start()
    for t in (h, w, s):
        t.join()


if __name__ == "__main__":
    main()