geMoldInsight/scripts/migrate_multi_upload.py

"""
数据库迁移脚本 - 添加多上传支持字段和验证字段

运行方式: python scripts/migrate_multi_upload.py

此脚本将:
1. 添加 upload_batch 字段到 stp_files 表
2. 添加 volume, surface_area, product_weight 快速查询字段到 stp_files 表
3. 移除 file_hash 字段的唯一约束（如果存在）
4. 为新字段创建索引
5. 添加验证结果字段到 analysis_metrics 表
"""
import asyncio
import sys
from pathlib import Path

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

from sqlalchemy import text
from src.database.database import db_manager
from src.utils.logger import get_logger

logger = get_logger(__name__)


async def check_column_exists(conn, table_name: str, column_name: str) -> bool:
    """检查列是否存在"""
    result = await conn.execute(text("""
        SELECT column_name 
        FROM information_schema.columns 
        WHERE table_name = :table_name 
        AND column_name = :column_name
    """), {"table_name": table_name, "column_name": column_name})
    return result.fetchone() is not None


async def check_index_exists(conn, index_name: str) -> bool:
    """检查索引是否存在"""
    result = await conn.execute(text("""
        SELECT indexname 
        FROM pg_indexes 
        WHERE indexname = :index_name
    """), {"index_name": index_name})
    return result.fetchone() is not None


async def run_migration():
    """执行迁移"""
    logger.info("开始数据库迁移 - 多上传支持和验证字段...")
    
    await db_manager.connect()
    
    if not db_manager.is_connected:
        logger.error("数据库连接失败")
        return False
    
    async with db_manager.engine.begin() as conn:
        migration_steps = []
        
        # stp_files 表字段
        if not await check_column_exists(conn, "stp_files", "upload_batch"):
            migration_steps.append("添加 upload_batch 字段")
            await conn.execute(text("""
                ALTER TABLE stp_files 
                ADD COLUMN upload_batch VARCHAR(36)
            """))
            logger.info("✓ 添加 upload_batch 字段")
        
        if not await check_column_exists(conn, "stp_files", "volume"):
            migration_steps.append("添加 volume 字段")
            await conn.execute(text("""
                ALTER TABLE stp_files 
                ADD COLUMN volume FLOAT
            """))
            logger.info("✓ 添加 volume 字段")
        
        if not await check_column_exists(conn, "stp_files", "surface_area"):
            migration_steps.append("添加 surface_area 字段")
            await conn.execute(text("""
                ALTER TABLE stp_files 
                ADD COLUMN surface_area FLOAT
            """))
            logger.info("✓ 添加 surface_area 字段")
        
        if not await check_column_exists(conn, "stp_files", "product_weight"):
            migration_steps.append("添加 product_weight 字段")
            await conn.execute(text("""
                ALTER TABLE stp_files 
                ADD COLUMN product_weight FLOAT
            """))
            logger.info("✓ 添加 product_weight 字段")
        
        # analysis_metrics 表验证字段
        if not await check_column_exists(conn, "analysis_metrics", "verification_status"):
            migration_steps.append("添加 verification_status 字段")
            await conn.execute(text("""
                ALTER TABLE analysis_metrics 
                ADD COLUMN verification_status VARCHAR(20)
            """))
            logger.info("✓ 添加 verification_status 字段")
        
        if not await check_column_exists(conn, "analysis_metrics", "verification_volume_diff"):
            migration_steps.append("添加 verification_volume_diff 字段")
            await conn.execute(text("""
                ALTER TABLE analysis_metrics 
                ADD COLUMN verification_volume_diff FLOAT
            """))
            logger.info("✓ 添加 verification_volume_diff 字段")
        
        if not await check_column_exists(conn, "analysis_metrics", "verification_area_diff"):
            migration_steps.append("添加 verification_area_diff 字段")
            await conn.execute(text("""
                ALTER TABLE analysis_metrics 
                ADD COLUMN verification_area_diff FLOAT
            """))
            logger.info("✓ 添加 verification_area_diff 字段")
        
        if not await check_column_exists(conn, "analysis_metrics", "verification_details"):
            migration_steps.append("添加 verification_details 字段")
            await conn.execute(text("""
                ALTER TABLE analysis_metrics 
                ADD COLUMN verification_details JSON
            """))
            logger.info("✓ 添加 verification_details 字段")
        
        # 索引
        if not await check_index_exists(conn, "ix_stp_files_upload_batch"):
            migration_steps.append("创建 upload_batch 索引")
            await conn.execute(text("""
                CREATE INDEX ix_stp_files_upload_batch 
                ON stp_files(upload_batch)
            """))
            logger.info("✓ 创建 upload_batch 索引")
        
        if not await check_index_exists(conn, "ix_stp_files_file_hash"):
            migration_steps.append("创建 file_hash 索引")
            await conn.execute(text("""
                CREATE INDEX ix_stp_files_file_hash 
                ON stp_files(file_hash)
            """))
            logger.info("✓ 创建 file_hash 索引")
        
        if not await check_index_exists(conn, "ix_stp_files_original_filename"):
            migration_steps.append("创建 original_filename 索引")
            await conn.execute(text("""
                CREATE INDEX ix_stp_files_original_filename 
                ON stp_files(original_filename)
            """))
            logger.info("✓ 创建 original_filename 索引")
        
        # 删除唯一约束
        try:
            result = await conn.execute(text("""
                SELECT conname 
                FROM pg_constraint 
                WHERE conrelid = 'stp_files'::regclass 
                AND contype = 'u'
            """))
            unique_constraints = result.fetchall()
            
            for constraint in unique_constraints:
                constraint_name = constraint[0]
                if 'file_hash' in constraint_name.lower():
                    migration_steps.append(f"删除唯一约束 {constraint_name}")
                    await conn.execute(text(f"""
                        ALTER TABLE stp_files 
                        DROP CONSTRAINT {constraint_name}
                    """))
                    logger.info(f"✓ 删除唯一约束: {constraint_name}")
        except Exception as e:
            logger.warning(f"检查唯一约束时出错（可能不存在）: {e}")
        
        if migration_steps:
            logger.info(f"\n迁移完成，执行了 {len(migration_steps)} 个步骤:")
            for step in migration_steps:
                logger.info(f"  - {step}")
        else:
            logger.info("\n无需迁移，所有字段和索引已存在")
    
    await db_manager.disconnect()
    logger.info("数据库迁移完成!")
    return True


async def rollback_migration():
    """回滚迁移"""
    logger.info("开始回滚数据库迁移...")
    
    await db_manager.connect()
    
    if not db_manager.is_connected:
        logger.error("数据库连接失败")
        return False
    
    async with db_manager.engine.begin() as conn:
        try:
            # stp_files 表
            if await check_index_exists(conn, "ix_stp_files_upload_batch"):
                await conn.execute(text("DROP INDEX IF EXISTS ix_stp_files_upload_batch"))
                logger.info("✓ 删除 upload_batch 索引")
            
            if await check_index_exists(conn, "ix_stp_files_file_hash"):
                await conn.execute(text("DROP INDEX IF EXISTS ix_stp_files_file_hash"))
                logger.info("✓ 删除 file_hash 索引")
            
            if await check_column_exists(conn, "stp_files", "upload_batch"):
                await conn.execute(text("ALTER TABLE stp_files DROP COLUMN upload_batch"))
                logger.info("✓ 删除 upload_batch 字段")
            
            if await check_column_exists(conn, "stp_files", "volume"):
                await conn.execute(text("ALTER TABLE stp_files DROP COLUMN volume"))
                logger.info("✓ 删除 volume 字段")
            
            if await check_column_exists(conn, "stp_files", "surface_area"):
                await conn.execute(text("ALTER TABLE stp_files DROP COLUMN surface_area"))
                logger.info("✓ 删除 surface_area 字段")
            
            if await check_column_exists(conn, "stp_files", "product_weight"):
                await conn.execute(text("ALTER TABLE stp_files DROP COLUMN product_weight"))
                logger.info("✓ 删除 product_weight 字段")
            
            # analysis_metrics 表
            if await check_column_exists(conn, "analysis_metrics", "verification_status"):
                await conn.execute(text("ALTER TABLE analysis_metrics DROP COLUMN verification_status"))
                logger.info("✓ 删除 verification_status 字段")
            
            if await check_column_exists(conn, "analysis_metrics", "verification_volume_diff"):
                await conn.execute(text("ALTER TABLE analysis_metrics DROP COLUMN verification_volume_diff"))
                logger.info("✓ 删除 verification_volume_diff 字段")
            
            if await check_column_exists(conn, "analysis_metrics", "verification_area_diff"):
                await conn.execute(text("ALTER TABLE analysis_metrics DROP COLUMN verification_area_diff"))
                logger.info("✓ 删除 verification_area_diff 字段")
            
            if await check_column_exists(conn, "analysis_metrics", "verification_details"):
                await conn.execute(text("ALTER TABLE analysis_metrics DROP COLUMN verification_details"))
                logger.info("✓ 删除 verification_details 字段")
            
            logger.info("回滚完成!")
        except Exception as e:
            logger.error(f"回滚失败: {e}")
            return False
    
    await db_manager.disconnect()
    return True


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="数据库迁移脚本 - 多上传支持和验证字段")
    parser.add_argument("--rollback", action="store_true", help="回滚迁移")
    args = parser.parse_args()
    
    if args.rollback:
        asyncio.run(rollback_migration())
    else:
        asyncio.run(run_migration())