From 224de8bbde2b578cc879f43a88f10c3281306ec3 Mon Sep 17 00:00:00 2001
From: SZCJW <792430652@qq.com>
Date: Sat, 7 Mar 2026 03:16:20 +0800
Subject: [PATCH] x

---
 scripts/migrate_multi_upload.py | 39 ++++++++++++++-------------------
 1 file changed, 16 insertions(+), 23 deletions(-)

diff --git a/scripts/migrate_multi_upload.py b/scripts/migrate_multi_upload.py
index 071a9c6..cb5c8ca 100644
--- a/scripts/migrate_multi_upload.py
+++ b/scripts/migrate_multi_upload.py
@@ -11,7 +11,6 @@
 """
 import asyncio
 import sys
-import os
 from pathlib import Path
 
 project_root = Path(__file__).parent.parent
@@ -19,7 +18,7 @@ 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.database.database import db_manager
 from src.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -46,31 +45,17 @@ async def check_index_exists(conn, index_name: str) -> bool:
     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
+    await db_manager.connect()
     
-    async with async_engine.begin() as conn:
+    if not db_manager.is_connected:
+        logger.error("数据库连接失败")
+        return False
+    
+    async with db_manager.engine.begin() as conn:
         migration_steps = []
         
         if not await check_column_exists(conn, "stp_files", "upload_batch"):
@@ -157,6 +142,7 @@ async def run_migration():
         else:
             logger.info("\n无需迁移，所有字段和索引已存在")
     
+    await db_manager.disconnect()
     logger.info("数据库迁移完成!")
     return True
 
@@ -165,7 +151,13 @@ async def rollback_migration():
     """回滚迁移"""
     logger.info("开始回滚数据库迁移...")
     
-    async with async_engine.begin() as conn:
+    await db_manager.connect()
+    
+    if not db_manager.is_connected:
+        logger.error("数据库连接失败")
+        return False
+    
+    async with db_manager.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"))
@@ -196,6 +188,7 @@ async def rollback_migration():
             logger.error(f"回滚失败: {e}")
             return False
     
+    await db_manager.disconnect()
     return True