"""
添加默认模型字段的迁移脚本
"""
from sqlalchemy import create_engine, text
from database import SQLALCHEMY_DATABASE_URL

def migrate():
    engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
    
    # 添加 is_default 字段
    with engine.connect() as conn:
        try:
            # 检查字段是否已存在
            result = conn.execute(text("PRAGMA table_info(ai_models)"))
            columns = [row[1] for row in result]
            
            if 'is_default' not in columns:
                print("添加 is_default 字段到 ai_models 表...")
                conn.execute(text("ALTER TABLE ai_models ADD COLUMN is_default BOOLEAN DEFAULT 0"))
                conn.commit()
                print("✓ is_default 字段添加成功")
            else:
                print("✓ is_default 字段已存在")
                
        except Exception as e:
            print(f"✗ 迁移失败: {e}")
            conn.rollback()
            raise

if __name__ == "__main__":
    print("开始数据库迁移...")
    migrate()
    print("迁移完成！")
