"""
添加工具代码字段的数据库迁移脚本
"""
import sqlite3

# 连接数据库
conn = sqlite3.connect('qa_system.db')
cursor = conn.cursor()

print("="*60)
print("数据库迁移：添加工具代码字段")
print("="*60)

try:
    # 检查字段是否已存在
    cursor.execute("PRAGMA table_info(ai_tools)")
    columns = [column[1] for column in cursor.fetchall()]
    
    if 'function_code' not in columns:
        print("添加 function_code 字段...")
        cursor.execute("ALTER TABLE ai_tools ADD COLUMN function_code TEXT")
        print("[OK] function_code 字段添加成功")
    else:
        print("[OK] function_code 字段已存在")
    
    if 'is_code_tool' not in columns:
        print("添加 is_code_tool 字段...")
        cursor.execute("ALTER TABLE ai_tools ADD COLUMN is_code_tool BOOLEAN DEFAULT 0")
        print("[OK] is_code_tool 字段添加成功")
    else:
        print("[OK] is_code_tool 字段已存在")
    
    conn.commit()
    print("\n" + "="*60)
    print("迁移完成！")
    print("="*60)
    
except Exception as e:
    print(f"[ERROR] 迁移失败: {e}")
    conn.rollback()

finally:
    conn.close()
