113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
直接使用 alembic API 执行数据库迁移的脚本
|
|
|
|
用法:
|
|
python scripts/direct_migration.py [--revision 版本号]
|
|
|
|
示例:
|
|
# 升级到最新版本
|
|
python scripts/direct_migration.py
|
|
|
|
# 升级到指定版本
|
|
python scripts/direct_migration.py --revision 001
|
|
|
|
# 降级到指定版本
|
|
python scripts/direct_migration.py --revision 001 --downgrade
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
import logging
|
|
|
|
# 配置日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.StreamHandler(sys.stdout),
|
|
logging.FileHandler('migration_direct.log', encoding='utf-8')
|
|
]
|
|
)
|
|
logger = logging.getLogger('alembic_direct')
|
|
|
|
# 将项目根目录添加到 Python 路径
|
|
root_dir = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(root_dir))
|
|
|
|
def run_migration(revision=None, downgrade=False, verbose=False):
|
|
"""
|
|
执行数据库迁移
|
|
|
|
Args:
|
|
revision (str): 版本号，为 None 表示升级到最新版本
|
|
downgrade (bool): 是否降级
|
|
verbose (bool): 是否显示详细日志
|
|
"""
|
|
try:
|
|
# 导入 alembic 相关模块
|
|
from alembic.config import Config
|
|
from alembic import command
|
|
|
|
# 设置日志级别
|
|
if verbose:
|
|
logging.getLogger('alembic').setLevel(logging.DEBUG)
|
|
|
|
# 获取配置文件路径
|
|
config_path = root_dir / "migrations" / "alembic.ini.tmp"
|
|
if not config_path.exists():
|
|
# 如果临时配置文件不存在，使用原始配置文件
|
|
config_path = root_dir / "migrations" / "alembic.ini"
|
|
|
|
logger.info(f"使用配置文件: {config_path}")
|
|
|
|
# 创建 alembic 配置对象
|
|
alembic_cfg = Config(str(config_path))
|
|
|
|
# 设置脚本位置
|
|
script_location = root_dir / "migrations"
|
|
alembic_cfg.set_main_option("script_location", str(script_location))
|
|
|
|
# 从 config/database.py 导入数据库 URL
|
|
from config.database import DATABASE_URL
|
|
logger.info(f"数据库 URL: {DATABASE_URL}")
|
|
|
|
# 设置数据库 URL
|
|
alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)
|
|
|
|
# 执行迁移
|
|
if downgrade:
|
|
logger.info(f"正在降级到版本: {revision or 'base'}")
|
|
command.downgrade(alembic_cfg, revision or "base")
|
|
else:
|
|
logger.info(f"正在升级到版本: {revision or 'head'}")
|
|
command.upgrade(alembic_cfg, revision or "head")
|
|
|
|
logger.info("迁移完成")
|
|
return True
|
|
except ImportError as e:
|
|
logger.error(f"导入错误: {e}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"执行迁移时发生错误: {e}", exc_info=True)
|
|
return False
|
|
|
|
|
|
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="显示详细日志")
|
|
|
|
args = parser.parse_args()
|
|
|
|
success = run_migration(args.revision, args.downgrade, args.verbose)
|
|
sys.exit(0 if success else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |