#!/usr/bin/env python3
"""
查找并杀掉数据库端的僵尸连接，释放 knowledge 表的锁
"""
import os, psycopg2
from dotenv import load_dotenv

_dir = os.path.dirname(os.path.abspath(__file__))
_root = os.path.normpath(os.path.join(_dir, '..', '..'))
load_dotenv(os.path.join(_root, '.env'))

conn = psycopg2.connect(
    host=os.getenv('KNOWHUB_DB'),
    port=int(os.getenv('KNOWHUB_PORT', 5432)),
    user=os.getenv('KNOWHUB_USER'),
    password=os.getenv('KNOWHUB_PASSWORD'),
    database=os.getenv('KNOWHUB_DB_NAME'),
    connect_timeout=10
)
conn.autocommit = True
cur = conn.cursor()

print("Connected.\n")

# 1. 查所有非 idle 的活跃连接
print("=== Active connections (non-idle) ===")
cur.execute("""
    SELECT pid, state, now() - query_start as duration, query
    FROM pg_stat_activity
    WHERE state != 'idle' AND pid != pg_backend_pid()
    ORDER BY query_start
""")
rows = cur.fetchall()
print(f"Found {len(rows)} active connections.\n")
for r in rows:
    print(f"  PID={r[0]}  state={r[1]}  duration={r[2]}")
    print(f"    query: {str(r[3])[:150]}")
    print()

# 2. 查所有 idle in transaction 的连接（这些最可能持有锁）
print("=== Idle-in-transaction connections ===")
cur.execute("""
    SELECT pid, state, now() - query_start as duration, query
    FROM pg_stat_activity
    WHERE state = 'idle in transaction' AND pid != pg_backend_pid()
    ORDER BY query_start
""")
idle_tx = cur.fetchall()
print(f"Found {len(idle_tx)} idle-in-transaction connections.\n")
for r in idle_tx:
    print(f"  PID={r[0]}  duration={r[2]}")
    print(f"    last query: {str(r[3])[:150]}")
    print()

# 3. 杀掉所有非当前的 非idle 和 idle-in-transaction 连接
cur.execute("""
    SELECT pid FROM pg_stat_activity
    WHERE pid != pg_backend_pid()
      AND (state != 'idle' OR state = 'idle in transaction')
""")
pids = [r[0] for r in cur.fetchall()]
print(f"=== Terminating {len(pids)} backend connections ===")
for pid in pids:
    try:
        cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
        print(f"  Killed PID {pid}")
    except Exception as e:
        print(f"  Failed PID {pid}: {e}")

print("\nDone. Try opening knowledge table in Navicat now.")
cur.close()
conn.close()