ai-LiaoWangweb-app/fix_pg_sequences.py

"""
修复 PostgreSQL 序列不同步问题。
当表中存在手动插入的数据时，自增序列会落后于实际最大 ID，
导致后续 INSERT 报 UniqueViolation。
"""
import os
from dotenv import load_dotenv

load_dotenv()

DB_HOST = os.environ.get("DB_HOST", "localhost")
DB_PORT = int(os.environ.get("DB_PORT", 5432))
DB_USER = os.environ.get("DB_USER", "")
DB_PASSWORD = os.environ.get("DB_PASSWORD", "")
DB_NAME = os.environ.get("DB_NAME", "")

TABLES = [
    "spider_task",        # collection_task
    "ai_conversation",
    "spider_result",
    "deep_collection",
    "ai_model",
    "token_usage_log",
    "ai_message",
    "user",
    "spider_source",
    "knowledge_import_task",
]


def fix_sequences():
    import psycopg2

    uri = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
    print(f"Connecting to {DB_HOST}:{DB_PORT}/{DB_NAME} as {DB_USER}...")

    conn = psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        dbname=DB_NAME,
    )
    conn.autocommit = True
    cur = conn.cursor()

    fixed = []
    skipped = []

    for table in TABLES:
        seq = f"{table}_id_seq"
        try:
            # 检查序列是否存在
            cur.execute(
                "SELECT EXISTS (SELECT 1 FROM information_schema.sequences WHERE sequence_name = %s)",
                (seq,),
            )
            exists = cur.fetchone()[0]
            if not exists:
                skipped.append(table)
                continue

            # 获取当前序列值和表最大 ID
            cur.execute(f"SELECT last_value FROM {seq}")
            seq_val = cur.fetchone()[0]

            cur.execute(f"SELECT COALESCE(MAX(id), 0) FROM {table}")
            max_id = cur.fetchone()[0]

            if max_id > seq_val:
                cur.execute(f"SELECT setval('{seq}', {max_id})")
                fixed.append(f"{table}: seq {seq_val} -> {max_id}")
            else:
                skipped.append(f"{table}: seq={seq_val}, max_id={max_id} (OK)")
        except Exception as e:
            skipped.append(f"{table}: ERROR - {e}")

    cur.close()
    conn.close()

    print("\n--- 修复结果 ---")
    for item in fixed:
        print(f"  [FIXED] {item}")
    for item in skipped:
        print(f"  [OK   ] {item}")

    if fixed:
        print(f"\n共修复 {len(fixed)} 个表的序列。")
    else:
        print("\n所有序列均正常，无需修复。")


if __name__ == "__main__":
    fix_sequences()