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

"""
数据库迁移脚本：为图像翻译表添加软删除字段

执行方式：
    python scripts/migrate_add_soft_delete.py

功能：
    1. 添加 is_deleted 字段（Boolean，默认 False）
    2. 添加 deleted_at 字段（DateTime，可为空）
    3. 创建 is_deleted 索引
"""
import sys
import os

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

from sqlalchemy import text
from app.database import engine, SessionLocal


def run_migration():
    """执行数据库迁移"""
    print("=" * 60)
    print("开始执行迁移：为图像翻译表添加软删除字段")
    print("=" * 60)
    
    db = SessionLocal()
    
    try:
        # 1. 添加 is_deleted 字段
        print("\n[1/4] 添加 is_deleted 字段...")
        try:
            db.execute(text("""
                ALTER TABLE aigcspace.image_translation 
                ADD COLUMN IF NOT EXISTS is_deleted BOOLEAN NOT NULL DEFAULT FALSE
            """))
            db.commit()
            print("✓ is_deleted 字段添加成功")
        except Exception as e:
            print(f"⚠ is_deleted 字段可能已存在: {e}")
            db.rollback()
        
        # 2. 添加 deleted_at 字段
        print("\n[2/4] 添加 deleted_at 字段...")
        try:
            db.execute(text("""
                ALTER TABLE aigcspace.image_translation 
                ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMP
            """))
            db.commit()
            print("✓ deleted_at 字段添加成功")
        except Exception as e:
            print(f"⚠ deleted_at 字段可能已存在: {e}")
            db.rollback()
        
        # 3. 创建索引
        print("\n[3/4] 创建 is_deleted 索引...")
        try:
            db.execute(text("""
                CREATE INDEX IF NOT EXISTS idx_image_translation_is_deleted 
                ON aigcspace.image_translation(is_deleted)
            """))
            db.commit()
            print("✓ 索引创建成功")
        except Exception as e:
            print(f"⚠ 索引可能已存在: {e}")
            db.rollback()
        
        # 4. 添加字段注释
        print("\n[4/4] 添加字段注释...")
        try:
            db.execute(text("""
                COMMENT ON COLUMN aigcspace.image_translation.is_deleted IS '是否已删除'
            """))
            db.execute(text("""
                COMMENT ON COLUMN aigcspace.image_translation.deleted_at IS '删除时间'
            """))
            db.commit()
            print("✓ 字段注释添加成功")
        except Exception as e:
            print(f"⚠ 添加注释失败: {e}")
            db.rollback()
        
        # 验证迁移结果
        print("\n" + "=" * 60)
        print("验证迁移结果...")
        print("=" * 60)
        
        result = db.execute(text("""
            SELECT 
                column_name, 
                data_type, 
                is_nullable,
                column_default
            FROM information_schema.columns
            WHERE table_schema = 'aigcspace'
            AND table_name = 'image_translation'
            AND column_name IN ('is_deleted', 'deleted_at')
            ORDER BY column_name
        """))
        
        columns = result.fetchall()
        if columns:
            print("\n新增字段信息：")
            for col in columns:
                print(f"  - {col[0]}: {col[1]} (nullable: {col[2]}, default: {col[3]})")
        else:
            print("⚠ 警告：未找到新增字段，请检查迁移是否成功")
        
        # 检查索引
        result = db.execute(text("""
            SELECT indexname
            FROM pg_indexes
            WHERE schemaname = 'aigcspace'
            AND tablename = 'image_translation'
            AND indexname = 'idx_image_translation_is_deleted'
        """))
        
        index = result.fetchone()
        if index:
            print(f"\n索引信息：")
            print(f"  - {index[0]}")
        else:
            print("\n⚠ 警告：未找到索引，请检查创建是否成功")
        
        # 统计现有记录
        result = db.execute(text("""
            SELECT COUNT(*) as total
            FROM aigcspace.image_translation
        """))
        total = result.fetchone()[0]
        
        print(f"\n当前表中记录总数: {total}")
        print(f"所有现有记录的 is_deleted 默认值为: FALSE")
        
        print("\n" + "=" * 60)
        print("✓ 迁移执行成功！")
        print("=" * 60)
        
        return True
        
    except Exception as e:
        print(f"\n✗ 迁移执行失败: {e}")
        import traceback
        traceback.print_exc()
        db.rollback()
        return False
        
    finally:
        db.close()


def rollback_migration():
    """回滚迁移（仅用于测试环境）"""
    print("\n" + "=" * 60)
    print("警告：准备回滚迁移")
    print("=" * 60)
    
    confirm = input("\n确定要回滚迁移吗？这将删除 is_deleted 和 deleted_at 字段！(yes/no): ")
    if confirm.lower() != 'yes':
        print("已取消回滚")
        return
    
    db = SessionLocal()
    
    try:
        # 删除索引
        print("\n[1/3] 删除索引...")
        db.execute(text("""
            DROP INDEX IF EXISTS aigcspace.idx_image_translation_is_deleted
        """))
        db.commit()
        print("✓ 索引删除成功")
        
        # 删除字段
        print("\n[2/3] 删除 is_deleted 字段...")
        db.execute(text("""
            ALTER TABLE aigcspace.image_translation 
            DROP COLUMN IF EXISTS is_deleted
        """))
        db.commit()
        print("✓ is_deleted 字段删除成功")
        
        print("\n[3/3] 删除 deleted_at 字段...")
        db.execute(text("""
            ALTER TABLE aigcspace.image_translation 
            DROP COLUMN IF EXISTS deleted_at
        """))
        db.commit()
        print("✓ deleted_at 字段删除成功")
        
        print("\n✓ 回滚完成")
        
    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(
        '--rollback',
        action='store_true',
        help='回滚迁移（仅用于测试环境）'
    )
    
    args = parser.parse_args()
    
    if args.rollback:
        rollback_migration()
    else:
        success = run_migration()
        sys.exit(0 if success else 1)