213 lines
8.1 KiB
Python
213 lines
8.1 KiB
Python
"""
|
|
数据库迁移脚本 - 添加多上传支持字段
|
|
|
|
运行方式: 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())
|