Agent/knowhub/knowhub_db/migrate_add_implemented_tools.py

#!/usr/bin/env python3
"""
数据库迁移脚本：为 tool_table 添加 implemented_tool_ids 字段

该字段用于存储工具库中实际接入的工具 ID 列表，建立工具描述与实际实现之间的映射关系。
"""

import os
import sys
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv

# 加载环境变量
_script_dir = os.path.dirname(os.path.abspath(__file__))
_project_root = os.path.normpath(os.path.join(_script_dir, '..', '..'))
load_dotenv(os.path.join(_project_root, '.env'))


def get_connection():
    """建立数据库连接"""
    host = os.getenv('KNOWHUB_DB')
    port = int(os.getenv('KNOWHUB_PORT', 5432))
    user = os.getenv('KNOWHUB_USER')
    password = os.getenv('KNOWHUB_PASSWORD')
    dbname = os.getenv('KNOWHUB_DB_NAME')

    print(f"连接到 {host}:{port}/{dbname} as {user} ...")
    conn = psycopg2.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=dbname,
        connect_timeout=10
    )
    conn.autocommit = True
    print("连接成功。")
    return conn


def main():
    print("=" * 60)
    print("开始迁移：添加 tool_table.implemented_tool_ids 字段")
    print("=" * 60)

    conn = get_connection()
    cursor = conn.cursor(cursor_factory=RealDictCursor)

    # 检查字段是否已存在
    cursor.execute("""
        SELECT column_name
        FROM information_schema.columns
        WHERE table_name = 'tool_table' AND column_name = 'implemented_tool_ids'
    """)

    if cursor.fetchone():
        print("\n字段 'implemented_tool_ids' 已存在，跳过迁移。")
    else:
        print("\n添加字段 'implemented_tool_ids' ...")
        cursor.execute("""
            ALTER TABLE tool_table
            ADD COLUMN implemented_tool_ids JSONB DEFAULT '[]'
        """)
        print("✓ 字段添加成功")

    # 验证结果
    print("\n" + "=" * 60)
    print("验证结果：")
    print("=" * 60)

    cursor.execute("""
        SELECT column_name, data_type, column_default
        FROM information_schema.columns
        WHERE table_name = 'tool_table'
        ORDER BY ordinal_position
    """)

    print("\ntool_table 当前字段：")
    for row in cursor.fetchall():
        print(f"  - {row['column_name']}: {row['data_type']} (默认: {row['column_default']})")

    cursor.execute("SELECT COUNT(*) as count FROM tool_table")
    count = cursor.fetchone()['count']
    print(f"\n总记录数: {count}")

    cursor.close()
    conn.close()
    print("\n迁移完成！")


if __name__ == '__main__':
    main()