Maas-Model-platform/backend/scripts/init_image_translation.py

"""
图像翻译模块数据库迁移脚本

功能：
1. 创建 image_translation 表
2. 插入模型配置（model 和 model_price）
3. 插入百宝箱应用数据（toolbox_apps）
4. 验证数据完整性

从 .env 文件读取数据库连接信息
"""
import os
import sys
from pathlib import Path
from dotenv import load_dotenv
import psycopg2
from psycopg2 import sql

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

# 加载环境变量
env_path = Path(__file__).parent.parent / '.env'
load_dotenv(dotenv_path=env_path)


def get_db_connection():
    """从环境变量获取数据库连接"""
    try:
        conn = psycopg2.connect(
            host=os.getenv("DB_HOST", "localhost"),
            port=int(os.getenv("DB_PORT", "5432")),
            database=os.getenv("DB_NAME", "aigcspace"),
            user=os.getenv("DB_USER", "postgres"),
            password=os.getenv("DB_PASSWORD", "")
        )
        return conn
    except Exception as e:
        print(f"❌ 数据库连接失败: {e}")
        sys.exit(1)


def execute_sql_file(cursor, file_path: Path):
    """执行SQL文件"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            sql_content = f.read()
        
        # 执行SQL
        cursor.execute(sql_content)
        print(f"✓ 执行成功: {file_path.name}")
        return True
    except Exception as e:
        print(f"❌ 执行失败: {file_path.name}")
        print(f"   错误信息: {e}")
        return False


def insert_model_data(cursor):
    """插入模型配置数据"""
    try:
        # 检查模型是否已存在
        cursor.execute("""
            SELECT id FROM aigcspace.model WHERE title = 'qwen-mt-image' LIMIT 1;
        """)
        existing_model = cursor.fetchone()
        
        if existing_model:
            print(f"✓ 模型 qwen-mt-image 已存在 (ID: {existing_model[0]})")
            model_id = existing_model[0]
        else:
            # 插入模型记录
            cursor.execute("""
                INSERT INTO aigcspace.model (
                    title, description, provider, model_type, 
                    status, created_at, updated_at
                ) VALUES (
                    'qwen-mt-image',
                    '通义千问图像翻译模型，支持多语种图像文字翻译',
                    'Alibaba',
                    'image-translation',
                    'active',
                    NOW(),
                    NOW()
                ) RETURNING id;
            """)
            model_id = cursor.fetchone()[0]
            print(f"✓ 模型 qwen-mt-image 插入成功 (ID: {model_id})")
        
        # 检查价格配置是否已存在
        cursor.execute("""
            SELECT id FROM aigcspace.model_price WHERE model_id = %s LIMIT 1;
        """, (model_id,))
        existing_price = cursor.fetchone()
        
        if existing_price:
            print(f"✓ 模型价格配置已存在 (ID: {existing_price[0]})")
        else:
            # 插入价格配置
            cursor.execute("""
                INSERT INTO aigcspace.model_price (
                    model_id, input_price, output_price, 
                    created_at, updated_at
                ) VALUES (
                    %s, 0.0000, 0.0200, NOW(), NOW()
                );
            """, (model_id,))
            print(f"✓ 模型价格配置插入成功（每张图片 0.02 元）")
        
        return True
    except Exception as e:
        print(f"❌ 模型数据插入失败: {e}")
        return False


def insert_toolbox_app(cursor):
    """插入百宝箱应用数据"""
    try:
        # 检查应用是否已存在
        cursor.execute("""
            SELECT id FROM aigcspace.toolbox_apps 
            WHERE app_key = 'qwen-image-translate' LIMIT 1;
        """)
        existing_app = cursor.fetchone()
        
        if existing_app:
            print(f"✓ 百宝箱应用 qwen-image-translate 已存在 (ID: {existing_app[0]})")
            return True
        
        # 插入应用记录
        cursor.execute("""
            INSERT INTO aigcspace.toolbox_apps (
                app_key, name, description, category, icon,
                is_new, is_hot, status, sort_order, version,
                created_at, updated_at
            ) VALUES (
                'qwen-image-translate',
                '通义千问图像翻译',
                '支持多语种图像文字翻译，保持原图排版和风格',
                '图像处理',
                'translate',
                true,
                true,
                'active',
                100,
                '1.0.0',
                NOW(),
                NOW()
            ) RETURNING id;
        """)
        app_id = cursor.fetchone()[0]
        print(f"✓ 百宝箱应用 qwen-image-translate 插入成功 (ID: {app_id})")
        
        return True
    except Exception as e:
        print(f"❌ 百宝箱应用插入失败: {e}")
        return False


def verify_data(cursor):
    """验证数据是否正确插入"""
    try:
        print(f"\n📊 数据验证:")
        
        # 验证表是否创建
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = 'aigcspace' 
                AND table_name = 'image_translation'
            );
        """)
        table_exists = cursor.fetchone()[0]
        print(f"   ✓ image_translation 表: {'已创建' if table_exists else '❌ 未创建'}")
        
        # 验证模型配置
        cursor.execute("""
            SELECT m.id, m.title, mp.output_price
            FROM aigcspace.model m
            LEFT JOIN aigcspace.model_price mp ON m.id = mp.model_id
            WHERE m.title = 'qwen-mt-image';
        """)
        model_data = cursor.fetchone()
        if model_data:
            print(f"   ✓ 模型配置: qwen-mt-image (ID: {model_data[0]}, 价格: {model_data[2]}元/张)")
        else:
            print(f"   ❌ 模型配置未找到")
        
        # 验证百宝箱应用
        cursor.execute("""
            SELECT id, name, category, is_new, version, status
            FROM aigcspace.toolbox_apps
            WHERE app_key = 'qwen-image-translate';
        """)
        app_data = cursor.fetchone()
        if app_data:
            app_id, name, category, is_new, version, status = app_data
            new_tag = " [New]" if is_new else ""
            print(f"   ✓ 百宝箱应用: {name}{new_tag} (版本: {version}, 状态: {status})")
        else:
            print(f"   ❌ 百宝箱应用未找到")
        
        # 验证索引
        cursor.execute("""
            SELECT indexname FROM pg_indexes 
            WHERE schemaname = 'aigcspace' 
            AND tablename = 'image_translation';
        """)
        indexes = cursor.fetchall()
        print(f"   ✓ 索引数量: {len(indexes)}")
        for idx in indexes:
            print(f"      - {idx[0]}")
        
        return True
    except Exception as e:
        print(f"❌ 数据验证失败: {e}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("图像翻译模块数据库迁移脚本")
    print("=" * 60)
    
    # 获取迁移文件路径
    migrations_dir = Path(__file__).parent.parent / 'migrations'
    migration_file = migrations_dir / 'add_image_translation_table.sql'
    
    # 检查文件是否存在
    if not migration_file.exists():
        print(f"❌ 迁移文件不存在: {migration_file}")
        sys.exit(1)
    
    print(f"\n📁 找到迁移文件:")
    print(f"   - {migration_file.name}")
    
    # 连接数据库
    print(f"\n🔌 连接数据库...")
    print(f"   Host: {os.getenv('DB_HOST', 'localhost')}")
    print(f"   Port: {os.getenv('DB_PORT', '5432')}")
    print(f"   Database: {os.getenv('DB_NAME', 'aigcspace')}")
    print(f"   User: {os.getenv('DB_USER', 'postgres')}")
    
    conn = get_db_connection()
    cursor = conn.cursor()
    
    print(f"✓ 数据库连接成功")
    
    # 执行迁移
    print(f"\n🚀 开始执行迁移...")
    
    # 1. 执行SQL文件（创建表）
    print(f"\n[1/3] 创建数据表...")
    if not execute_sql_file(cursor, migration_file):
        conn.rollback()
        cursor.close()
        conn.close()
        print(f"\n❌ 迁移失败，已回滚所有更改")
        sys.exit(1)
    
    
    # 提交事务
    conn.commit()
    print(f"\n✓ 所有迁移执行成功")
    
    # 验证数据
    verify_data(cursor)
    
    # 关闭连接
    cursor.close()
    conn.close()
    
    print(f"\n" + "=" * 60)
    print("✅ 图像翻译模块数据库迁移完成！")
    print("=" * 60)
    print(f"\n💡 下一步:")
    print(f"   1. 重启后端服务: cd backend && python main.py")
    print(f"   2. 访问百宝箱: http://localhost:5173/toolbox")
    print(f"   3. 点击「通义千问图像翻译」卡片的「立即使用」按钮")
    print(f"   4. 开始使用图像翻译功能")
    print()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print(f"\n\n⚠️  用户中断执行")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ 执行过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)