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

"""
百宝箱数据库迁移脚本

执行百宝箱相关的SQL迁移文件
从.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 verify_data(cursor):
    """验证数据是否正确插入"""
    try:
        # 查询总数
        cursor.execute("SELECT COUNT(*) FROM aigcspace.toolbox_apps;")
        total = cursor.fetchone()[0]
        print(f"\n📊 数据验证:")
        print(f"   总应用数: {total}")
        
        # 按分类统计
        cursor.execute("""
            SELECT category, COUNT(*) as count
            FROM aigcspace.toolbox_apps
            WHERE status = 'active'
            GROUP BY category
            ORDER BY count DESC;
        """)
        
        print(f"\n   分类分布:")
        for row in cursor.fetchall():
            print(f"   - {row[0]}: {row[1]}个")
        
        # 显示所有应用
        cursor.execute("""
            SELECT app_key, name, category, is_new, version
            FROM aigcspace.toolbox_apps
            ORDER BY sort_order DESC;
        """)
        
        print(f"\n   应用列表:")
        for row in cursor.fetchall():
            app_key, name, category, is_new, version = row
            new_tag = " [New]" if is_new else ""
            version_tag = f" ({version})" if version else ""
            print(f"   - {name}{new_tag}{version_tag} - {category}")
        
        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_files = [
        migrations_dir / '034_create_toolbox_apps_table.sql',
        migrations_dir / '035_init_toolbox_apps_data.sql'
    ]
    
    # 检查文件是否存在
    for file_path in migration_files:
        if not file_path.exists():
            print(f"❌ 迁移文件不存在: {file_path}")
            sys.exit(1)
    
    print(f"\n📁 找到 {len(migration_files)} 个迁移文件")
    for file_path in migration_files:
        print(f"   - {file_path.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🚀 开始执行迁移...")
    success_count = 0
    
    for file_path in migration_files:
        if execute_sql_file(cursor, file_path):
            success_count += 1
        else:
            # 如果执行失败，回滚并退出
            conn.rollback()
            cursor.close()
            conn.close()
            print(f"\n❌ 迁移失败，已回滚所有更改")
            sys.exit(1)
    
    # 提交事务
    conn.commit()
    print(f"\n✓ 所有迁移执行成功 ({success_count}/{len(migration_files)})")
    
    # 验证数据
    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. 启动前端服务: cd frontend && npm run dev")
    print(f"   3. 访问页面: http://localhost:5173/toolbox")
    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)