341 lines
11 KiB
Python
341 lines
11 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
数据库迁移执行脚本
|
|
|
|
用法:
|
|
python scripts/run_migration.py [--revision 版本号]
|
|
|
|
示例:
|
|
# 升级到最新版本
|
|
python scripts/run_migration.py
|
|
|
|
# 升级到指定版本
|
|
python scripts/run_migration.py --revision 001
|
|
|
|
# 降级到指定版本
|
|
python scripts/run_migration.py --revision 001 --downgrade
|
|
|
|
# 生成新的迁移脚本
|
|
python scripts/run_migration.py --generate "添加新字段"
|
|
|
|
# 为特定表生成迁移脚本
|
|
python scripts/run_migration.py --generate "为用户表添加邮箱字段" --table users
|
|
|
|
# 为多个表生成迁移脚本
|
|
python scripts/run_migration.py --generate "添加审计字段" --table users,orders,products
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
import subprocess
|
|
import shutil
|
|
import tempfile
|
|
import locale
|
|
import datetime
|
|
|
|
# 将项目根目录添加到 Python 路径
|
|
root_dir = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(root_dir))
|
|
|
|
# 导入项目日志模块
|
|
from utils.logger import get_logger, setup_logger
|
|
|
|
# 设置日志
|
|
setup_logger()
|
|
logger = get_logger('migration')
|
|
|
|
def create_ascii_config(original_config_path):
|
|
"""
|
|
创建一个 ASCII 编码的配置文件副本
|
|
|
|
Args:
|
|
original_config_path (Path): 原始配置文件路径
|
|
|
|
Returns:
|
|
Path: 临时配置文件路径
|
|
"""
|
|
# 如果临时文件已存在，直接使用
|
|
temp_config_path = original_config_path.with_suffix('.ini.tmp')
|
|
if temp_config_path.exists():
|
|
logger.info(f"使用已存在的临时配置文件: {temp_config_path}")
|
|
return temp_config_path
|
|
|
|
# 创建临时文件
|
|
logger.info(f"创建 ASCII 编码的临时配置文件: {temp_config_path}")
|
|
|
|
# 复制配置文件内容，去除中文注释
|
|
with open(temp_config_path, 'w', encoding='ascii', errors='ignore') as temp_file:
|
|
temp_file.write("""
|
|
# Alembic Configuration File
|
|
# This file contains basic configuration for Alembic
|
|
|
|
[alembic]
|
|
# Path to migration scripts
|
|
script_location = migrations
|
|
|
|
# Template uses jinja2 format
|
|
output_encoding = utf-8
|
|
|
|
# Database connection configuration
|
|
# In practice, this value will be overridden by the configuration in env.py
|
|
sqlalchemy.url = driver://user:pass@localhost/dbname
|
|
|
|
# Logging configuration
|
|
[loggers]
|
|
keys = root,sqlalchemy,alembic
|
|
|
|
[handlers]
|
|
keys = console
|
|
|
|
[formatters]
|
|
keys = generic
|
|
|
|
[logger_root]
|
|
level = WARN
|
|
handlers = console
|
|
qualname =
|
|
|
|
[logger_sqlalchemy]
|
|
level = WARN
|
|
handlers =
|
|
qualname = sqlalchemy.engine
|
|
|
|
[logger_alembic]
|
|
level = INFO
|
|
handlers =
|
|
qualname = alembic
|
|
|
|
[handler_console]
|
|
class = StreamHandler
|
|
args = (sys.stderr,)
|
|
level = NOTSET
|
|
formatter = generic
|
|
|
|
[formatter_generic]
|
|
format = %(levelname)-5.5s [%(name)s] %(message)s
|
|
datefmt = %H:%M:%S
|
|
""")
|
|
|
|
return temp_config_path
|
|
|
|
def run_migration_command(command, args):
|
|
"""
|
|
执行迁移命令
|
|
|
|
Args:
|
|
command (str): 命令名称，如 'upgrade', 'downgrade', 'revision'
|
|
args (argparse.Namespace): 命令行参数
|
|
|
|
Returns:
|
|
bool: 是否成功执行命令
|
|
"""
|
|
# 获取项目根目录
|
|
root_dir = Path(__file__).parent.parent
|
|
|
|
# 原始配置文件路径
|
|
original_config_path = root_dir / "migrations" / "alembic.ini"
|
|
|
|
# 创建 ASCII 编码的临时配置文件
|
|
temp_config_path = create_ascii_config(original_config_path)
|
|
|
|
# 构建命令
|
|
cmd = ["alembic", "-c", str(temp_config_path)]
|
|
|
|
if args.verbose:
|
|
cmd.append("--verbose")
|
|
|
|
# 打印所有参数，帮助调试
|
|
logger.info(f"命令参数: {vars(args)}")
|
|
|
|
if command == 'upgrade':
|
|
cmd.extend(["upgrade", args.revision or "head"])
|
|
elif command == 'downgrade':
|
|
cmd.extend(["downgrade", args.revision or "base"])
|
|
elif command == 'revision':
|
|
logger.info(f"生成迁移脚本，描述: {args.message}")
|
|
cmd.extend(["revision", "--autogenerate", "-m", args.message])
|
|
|
|
# 如果指定了表名，创建一个临时的 env.py 文件，只包含指定的表
|
|
if args.table:
|
|
# 将表名列表转换为 Python 列表字符串
|
|
tables = [f"'{table.strip()}'" for table in args.table.split(',')]
|
|
tables_str = f"[{', '.join(tables)}]"
|
|
|
|
# 创建临时环境变量，传递表名列表
|
|
os.environ["ALEMBIC_TABLES"] = args.table
|
|
logger.info(f"指定迁移表: {args.table}")
|
|
|
|
if args.branch:
|
|
cmd.extend(["--branch", args.branch])
|
|
elif command == 'history':
|
|
cmd.append("history")
|
|
if args.verbose:
|
|
cmd.append("-v")
|
|
elif command == 'current':
|
|
cmd.append("current")
|
|
elif command == 'show':
|
|
cmd.extend(["show", args.revision or "head"])
|
|
elif command == 'list_tables':
|
|
# 这不是 alembic 命令，而是我们自定义的命令
|
|
return list_database_tables()
|
|
else:
|
|
logger.error(f"未知命令: {command}")
|
|
return False
|
|
|
|
# 执行命令
|
|
logger.info(f"执行命令: {' '.join(cmd)}")
|
|
logger.info(f"工作目录: {root_dir}")
|
|
logger.info(f"Python 编码: {sys.getdefaultencoding()}")
|
|
logger.info(f"文件系统编码: {sys.getfilesystemencoding()}")
|
|
logger.info(f"系统默认编码: {locale.getpreferredencoding()}")
|
|
|
|
# 设置环境变量，确保使用 UTF-8 编码
|
|
env = os.environ.copy()
|
|
env["PYTHONIOENCODING"] = "utf-8"
|
|
|
|
# 创建命令输出日志文件 - 使用项目日志目录
|
|
from config.settings import LogConfig
|
|
log_config = LogConfig.as_dict()
|
|
log_dir = Path(os.path.dirname(log_config["file"]))
|
|
log_dir.mkdir(exist_ok=True)
|
|
|
|
# 使用日期和命令类型命名日志文件
|
|
today = datetime.datetime.now().strftime('%Y%m%d')
|
|
log_file = log_dir / f"migration_{command}_{today}.log"
|
|
|
|
try:
|
|
# 使用直接执行命令的方式，避免 subprocess 的编码问题
|
|
logger.info(f"正在执行迁移，日志将写入 {log_file}...")
|
|
|
|
# 使用 subprocess.Popen 而不是 subprocess.run
|
|
with open(log_file, "a", encoding="utf-8") as f_out:
|
|
# 添加分隔线和时间戳
|
|
f_out.write(f"\n\n{'='*80}\n")
|
|
f_out.write(f"执行时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
|
f_out.write(f"执行命令: {' '.join(cmd)}\n")
|
|
f_out.write(f"{'='*80}\n\n")
|
|
|
|
process = subprocess.Popen(
|
|
cmd,
|
|
cwd=str(root_dir),
|
|
env=env,
|
|
stdout=f_out,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
encoding="utf-8",
|
|
errors="replace" # 使用 replace 策略处理无法解码的字符
|
|
)
|
|
process.wait()
|
|
|
|
# 读取日志文件
|
|
with open(log_file, "r", encoding="utf-8") as f:
|
|
# 只读取最后 50 行，避免日志过长
|
|
lines = f.readlines()
|
|
output = ''.join(lines[-50:]) if len(lines) > 50 else ''.join(lines)
|
|
|
|
logger.info(f"命令执行状态: {'成功' if process.returncode == 0 else '失败'}")
|
|
logger.info(f"输出(最后部分):\n{output}")
|
|
|
|
# 清理环境变量
|
|
if 'ALEMBIC_TABLES' in os.environ:
|
|
del os.environ['ALEMBIC_TABLES']
|
|
|
|
return process.returncode == 0
|
|
except Exception as e:
|
|
logger.error(f"执行命令时发生错误: {e}", exc_info=True)
|
|
|
|
# 清理环境变量
|
|
if 'ALEMBIC_TABLES' in os.environ:
|
|
del os.environ['ALEMBIC_TABLES']
|
|
|
|
return False
|
|
|
|
def list_database_tables():
|
|
"""
|
|
列出数据库中的所有表
|
|
|
|
Returns:
|
|
bool: 是否成功执行命令
|
|
"""
|
|
try:
|
|
# 导入数据库配置
|
|
from config.database import DBConfig
|
|
|
|
# 获取数据库连接
|
|
engine = DBConfig.engine
|
|
|
|
# 获取所有表名
|
|
from sqlalchemy import inspect
|
|
inspector = inspect(engine)
|
|
tables = inspector.get_table_names()
|
|
|
|
logger.info("数据库中的表:")
|
|
for i, table in enumerate(tables, 1):
|
|
logger.info(f"{i}. {table}")
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"获取数据库表时发生错误: {e}", exc_info=True)
|
|
return False
|
|
|
|
def run_migration(args):
|
|
"""
|
|
执行数据库迁移
|
|
|
|
Args:
|
|
args (argparse.Namespace): 命令行参数
|
|
|
|
Returns:
|
|
bool: 是否成功执行迁移
|
|
"""
|
|
# 打印所有参数，帮助调试
|
|
logger.info(f"运行迁移，参数: {vars(args)}")
|
|
|
|
if args.list_tables:
|
|
# 列出数据库中的所有表
|
|
return run_migration_command('list_tables', args)
|
|
elif args.message:
|
|
# 生成新的迁移脚本
|
|
logger.info(f"生成新的迁移脚本，描述: {args.message}")
|
|
return run_migration_command('revision', args)
|
|
elif args.history:
|
|
# 显示迁移历史
|
|
return run_migration_command('history', args)
|
|
elif args.current:
|
|
# 显示当前版本
|
|
return run_migration_command('current', args)
|
|
elif args.show:
|
|
# 显示指定版本的详细信息
|
|
return run_migration_command('show', args)
|
|
elif args.downgrade:
|
|
# 降级到指定版本
|
|
return run_migration_command('downgrade', args)
|
|
else:
|
|
# 升级到指定版本
|
|
return run_migration_command('upgrade', args)
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="执行数据库迁移")
|
|
parser.add_argument("--revision", help="版本号，为空表示升级到最新版本")
|
|
parser.add_argument("--downgrade", action="store_true", help="是否降级")
|
|
parser.add_argument("--verbose", "-v", action="store_true", help="显示详细日志")
|
|
parser.add_argument("--generate", "--gen", "-m", dest="message", help="生成新的迁移脚本，需要提供描述信息")
|
|
parser.add_argument("--branch", "-b", help="分支名称，用于生成迁移脚本")
|
|
parser.add_argument("--history", action="store_true", help="显示迁移历史")
|
|
parser.add_argument("--current", action="store_true", help="显示当前版本")
|
|
parser.add_argument("--show", action="store_true", help="显示指定版本的详细信息")
|
|
parser.add_argument("--table", "-t", help="指定要迁移的表名，多个表用逗号分隔")
|
|
parser.add_argument("--list-tables", action="store_true", help="列出数据库中的所有表")
|
|
|
|
args = parser.parse_args()
|
|
|
|
success = run_migration(args)
|
|
sys.exit(0 if success else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |