#!/usr/bin/env python3
"""数据库迁移脚本：为algorithm_services表添加技术分类和输出类型字段"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from sqlalchemy import text
from app.models.database import engine

def migrate():
    """执行数据库迁移"""
    try:
        with engine.connect() as conn:
            # 检查字段是否已存在（PostgreSQL语法）
            result = conn.execute(text("""
                SELECT column_name 
                FROM information_schema.columns 
                WHERE table_name = 'algorithm_services'
            """))
            columns = [row[0] for row in result.fetchall()]
            
            # 添加 tech_category 字段
            if 'tech_category' not in columns:
                conn.execute(text("ALTER TABLE algorithm_services ADD COLUMN tech_category VARCHAR(50) DEFAULT 'computer_vision'"))
                print("✓ 已添加 tech_category 字段到 algorithm_services 表")
            else:
                print("✓ tech_category 字段已存在于 algorithm_services 表")
            
            # 添加 output_type 字段
            if 'output_type' not in columns:
                conn.execute(text("ALTER TABLE algorithm_services ADD COLUMN output_type VARCHAR(50) DEFAULT 'image'"))
                print("✓ 已添加 output_type 字段到 algorithm_services 表")
            else:
                print("✓ output_type 字段已存在于 algorithm_services 表")
            
            conn.commit()
            print("\n数据库迁移完成！")
            
    except Exception as e:
        print(f"数据库迁移失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    migrate()