""" 数据库迁移脚本 - 添加多上传支持字段 运行方式: python scripts/migrate_multi_upload.py 此脚本将: 1. 添加 upload_batch 字段到 stp_files 表 2. 添加 volume, surface_area, product_weight 快速查询字段到 stp_files 表 3. 移除 file_hash 字段的唯一约束（如果存在） 4. 为新字段创建索引 """ import asyncio import sys import os 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 async_engine 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 check_unique_constraint_exists(conn, table_name: str, column_name: str) -> bool: """检查唯一约束是否存在""" result = await conn.execute(text(""" SELECT conname FROM pg_constraint WHERE conrelid = :table_name::regclass AND contype = 'u' AND conname LIKE :pattern """), {"table_name": table_name, "pattern": f"%{column_name}%"}) return result.fetchone() is not None async def run_migration(): """执行迁移""" logger.info("开始数据库迁移 - 多上传支持...") async with async_engine.begin() as conn: try: await conn.execute(text("SELECT 1")) logger.info("数据库连接成功") except Exception as e: logger.error(f"数据库连接失败: {e}") return False async with async_engine.begin() as conn: migration_steps = [] 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 字段") 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无需迁移，所有字段和索引已存在") logger.info("数据库迁移完成!") return True async def rollback_migration(): """回滚迁移""" logger.info("开始回滚数据库迁移...") async with async_engine.begin() as conn: try: 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 字段") logger.info("回滚完成!") except Exception as e: logger.error(f"回滚失败: {e}") return False 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())