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

"""
数据库迁移管理脚本

执行方式：
    # 执行所有待执行的迁移
    python scripts/run_migrations.py
    
    # 执行特定迁移
    python scripts/run_migrations.py --migration migrate_add_soft_delete
    
    # 查看迁移状态
    python scripts/run_migrations.py --status
    
    # 回滚最后一次迁移
    python scripts/run_migrations.py --rollback
"""
import sys
import os
from datetime import datetime

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy import text
from app.database import SessionLocal


# 迁移列表（按执行顺序）
MIGRATIONS = [
    {
        'name': 'migrate_add_soft_delete',
        'description': '为图像翻译表添加软删除字段',
        'module': 'scripts.migrate_add_soft_delete',
        'version': '20260122_001'
    },
    # 在这里添加新的迁移...
]


def ensure_migration_table():
    """确保迁移记录表存在"""
    db = SessionLocal()
    try:
        db.execute(text("""
            CREATE TABLE IF NOT EXISTS aigcspace.schema_migrations (
                id SERIAL PRIMARY KEY,
                version VARCHAR(50) UNIQUE NOT NULL,
                name VARCHAR(255) NOT NULL,
                description TEXT,
                executed_at TIMESTAMP NOT NULL DEFAULT NOW(),
                execution_time_ms INTEGER,
                status VARCHAR(20) NOT NULL DEFAULT 'success'
            )
        """))
        db.commit()
        print("✓ 迁移记录表已就绪")
    except Exception as e:
        print(f"✗ 创建迁移记录表失败: {e}")
        db.rollback()
        raise
    finally:
        db.close()


def is_migration_executed(version):
    """检查迁移是否已执行"""
    db = SessionLocal()
    try:
        result = db.execute(
            text("SELECT COUNT(*) FROM aigcspace.schema_migrations WHERE version = :version"),
            {"version": version}
        )
        count = result.scalar()
        return count > 0
    finally:
        db.close()


def record_migration(version, name, description, execution_time_ms, status='success'):
    """记录迁移执行"""
    db = SessionLocal()
    try:
        db.execute(
            text("""
                INSERT INTO aigcspace.schema_migrations 
                (version, name, description, execution_time_ms, status)
                VALUES (:version, :name, :description, :execution_time_ms, :status)
            """),
            {
                "version": version,
                "name": name,
                "description": description,
                "execution_time_ms": execution_time_ms,
                "status": status
            }
        )
        db.commit()
    except Exception as e:
        print(f"✗ 记录迁移失败: {e}")
        db.rollback()
    finally:
        db.close()


def run_migration(migration_info):
    """执行单个迁移"""
    version = migration_info['version']
    name = migration_info['name']
    description = migration_info['description']
    module_name = migration_info['module']
    
    print("\n" + "=" * 70)
    print(f"执行迁移: {name}")
    print(f"版本: {version}")
    print(f"描述: {description}")
    print("=" * 70)
    
    # 检查是否已执行
    if is_migration_executed(version):
        print(f"⚠ 迁移 {version} 已执行，跳过")
        return True
    
    try:
        # 动态导入迁移模块
        module = __import__(module_name, fromlist=['run_migration'])
        
        # 执行迁移
        start_time = datetime.now()
        success = module.run_migration()
        end_time = datetime.now()
        
        execution_time_ms = int((end_time - start_time).total_seconds() * 1000)
        
        if success:
            # 记录成功的迁移
            record_migration(version, name, description, execution_time_ms, 'success')
            print(f"\n✓ 迁移 {name} 执行成功 (耗时: {execution_time_ms}ms)")
            return True
        else:
            # 记录失败的迁移
            record_migration(version, name, description, execution_time_ms, 'failed')
            print(f"\n✗ 迁移 {name} 执行失败")
            return False
            
    except Exception as e:
        print(f"\n✗ 迁移 {name} 执行异常: {e}")
        import traceback
        traceback.print_exc()
        
        # 记录失败
        record_migration(version, name, description, 0, 'error')
        return False


def run_all_migrations():
    """执行所有待执行的迁移"""
    print("\n" + "=" * 70)
    print("开始执行数据库迁移")
    print("=" * 70)
    
    # 确保迁移表存在
    ensure_migration_table()
    
    # 统计
    total = len(MIGRATIONS)
    executed = 0
    skipped = 0
    failed = 0
    
    # 执行每个迁移
    for migration in MIGRATIONS:
        if is_migration_executed(migration['version']):
            skipped += 1
            print(f"\n⚠ 跳过已执行的迁移: {migration['name']} ({migration['version']})")
            continue
        
        success = run_migration(migration)
        if success:
            executed += 1
        else:
            failed += 1
            print(f"\n✗ 迁移失败，停止执行后续迁移")
            break
    
    # 输出总结
    print("\n" + "=" * 70)
    print("迁移执行总结")
    print("=" * 70)
    print(f"总迁移数: {total}")
    print(f"已执行: {executed}")
    print(f"已跳过: {skipped}")
    print(f"失败: {failed}")
    print("=" * 70)
    
    return failed == 0


def show_migration_status():
    """显示迁移状态"""
    print("\n" + "=" * 70)
    print("数据库迁移状态")
    print("=" * 70)
    
    # 确保迁移表存在
    ensure_migration_table()
    
    db = SessionLocal()
    try:
        # 查询已执行的迁移
        result = db.execute(text("""
            SELECT version, name, description, executed_at, execution_time_ms, status
            FROM aigcspace.schema_migrations
            ORDER BY executed_at DESC
        """))
        
        executed_migrations = {row[0]: row for row in result.fetchall()}
        
        print(f"\n待执行的迁移:")
        print("-" * 70)
        
        pending_count = 0
        for migration in MIGRATIONS:
            version = migration['version']
            if version not in executed_migrations:
                pending_count += 1
                print(f"  [{version}] {migration['name']}")
                print(f"    描述: {migration['description']}")
                print()
        
        if pending_count == 0:
            print("  无待执行的迁移")
        
        print(f"\n已执行的迁移:")
        print("-" * 70)
        
        if executed_migrations:
            for migration in MIGRATIONS:
                version = migration['version']
                if version in executed_migrations:
                    row = executed_migrations[version]
                    status_icon = "✓" if row[5] == 'success' else "✗"
                    print(f"  {status_icon} [{version}] {row[1]}")
                    print(f"    执行时间: {row[3]}")
                    print(f"    耗时: {row[4]}ms")
                    print(f"    状态: {row[5]}")
                    print()
        else:
            print("  无已执行的迁移")
        
        print("=" * 70)
        
    finally:
        db.close()


def rollback_last_migration():
    """回滚最后一次迁移"""
    print("\n" + "=" * 70)
    print("警告：准备回滚最后一次迁移")
    print("=" * 70)
    
    db = SessionLocal()
    try:
        # 查询最后一次迁移
        result = db.execute(text("""
            SELECT version, name, description
            FROM aigcspace.schema_migrations
            ORDER BY executed_at DESC
            LIMIT 1
        """))
        
        last_migration = result.fetchone()
        if not last_migration:
            print("\n没有可回滚的迁移")
            return
        
        version, name, description = last_migration
        print(f"\n最后一次迁移:")
        print(f"  版本: {version}")
        print(f"  名称: {name}")
        print(f"  描述: {description}")
        
        confirm = input("\n确定要回滚此迁移吗？(yes/no): ")
        if confirm.lower() != 'yes':
            print("已取消回滚")
            return
        
        # 查找对应的迁移配置
        migration_info = None
        for m in MIGRATIONS:
            if m['version'] == version:
                migration_info = m
                break
        
        if not migration_info:
            print(f"\n✗ 未找到迁移配置: {version}")
            return
        
        # 动态导入迁移模块
        module_name = migration_info['module']
        module = __import__(module_name, fromlist=['rollback_migration'])
        
        if not hasattr(module, 'rollback_migration'):
            print(f"\n✗ 迁移 {name} 不支持回滚")
            return
        
        # 执行回滚
        module.rollback_migration()
        
        # 删除迁移记录
        db.execute(
            text("DELETE FROM aigcspace.schema_migrations WHERE version = :version"),
            {"version": version}
        )
        db.commit()
        
        print(f"\n✓ 迁移 {name} 已回滚")
        
    except Exception as e:
        print(f"\n✗ 回滚失败: {e}")
        import traceback
        traceback.print_exc()
        db.rollback()
        
    finally:
        db.close()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='数据库迁移管理工具')
    parser.add_argument(
        '--migration',
        type=str,
        help='执行指定的迁移'
    )
    parser.add_argument(
        '--status',
        action='store_true',
        help='显示迁移状态'
    )
    parser.add_argument(
        '--rollback',
        action='store_true',
        help='回滚最后一次迁移'
    )
    
    args = parser.parse_args()
    
    try:
        if args.status:
            show_migration_status()
        elif args.rollback:
            rollback_last_migration()
        elif args.migration:
            # 执行指定迁移
            migration_info = None
            for m in MIGRATIONS:
                if m['name'] == args.migration:
                    migration_info = m
                    break
            
            if migration_info:
                ensure_migration_table()
                success = run_migration(migration_info)
                sys.exit(0 if success else 1)
            else:
                print(f"✗ 未找到迁移: {args.migration}")
                sys.exit(1)
        else:
            # 执行所有待执行的迁移
            success = run_all_migrations()
            sys.exit(0 if success else 1)
            
    except KeyboardInterrupt:
        print("\n\n用户中断执行")
        sys.exit(1)
    except Exception as e:
        print(f"\n✗ 执行失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)