xiamuceer
184 lines
5.5 KiB
Python
184 lines
5.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
数据库自动迁移脚本
|
|
用于开发和生产环境的数据库迁移管理
|
|
"""
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# 添加项目路径
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from app.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def run_command(cmd: list, description: str) -> bool:
|
|
"""运行命令并返回是否成功"""
|
|
try:
|
|
logger.info(f"🚀 {description}...")
|
|
result = subprocess.run(
|
|
cmd,
|
|
cwd=project_root,
|
|
capture_output=True,
|
|
text=True,
|
|
check=False
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
logger.info(f"✅ {description}成功")
|
|
if result.stdout:
|
|
print(result.stdout)
|
|
return True
|
|
else:
|
|
logger.error(f"❌ {description}失败")
|
|
if result.stderr:
|
|
print(result.stderr, file=sys.stderr)
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"❌ {description}异常: {e}")
|
|
return False
|
|
|
|
|
|
def create_migration(message: str = None):
|
|
"""创建新的迁移版本"""
|
|
if not message:
|
|
message = input("请输入迁移描述: ").strip()
|
|
if not message:
|
|
message = "auto_migration"
|
|
|
|
cmd = ["alembic", "revision", "--autogenerate", "-m", message]
|
|
return run_command(cmd, f"生成迁移: {message}")
|
|
|
|
|
|
def upgrade_database(revision: str = "head"):
|
|
"""升级数据库到指定版本"""
|
|
cmd = ["alembic", "upgrade", revision]
|
|
return run_command(cmd, f"升级数据库到: {revision}")
|
|
|
|
|
|
def downgrade_database(revision: str = "-1"):
|
|
"""降级数据库到指定版本"""
|
|
cmd = ["alembic", "downgrade", revision]
|
|
return run_command(cmd, f"降级数据库到: {revision}")
|
|
|
|
|
|
def show_current():
|
|
"""显示当前数据库版本"""
|
|
cmd = ["alembic", "current"]
|
|
return run_command(cmd, "查看当前版本")
|
|
|
|
|
|
def show_history():
|
|
"""显示迁移历史"""
|
|
cmd = ["alembic", "history", "--verbose"]
|
|
return run_command(cmd, "查看迁移历史")
|
|
|
|
|
|
def show_heads():
|
|
"""显示最新版本"""
|
|
cmd = ["alembic", "heads"]
|
|
return run_command(cmd, "查看最新版本")
|
|
|
|
|
|
def stamp_database(revision: str = "head"):
|
|
"""标记数据库版本（不执行迁移）"""
|
|
cmd = ["alembic", "stamp", revision]
|
|
return run_command(cmd, f"标记数据库版本: {revision}")
|
|
|
|
|
|
def auto_migrate():
|
|
"""自动迁移：生成并执行迁移"""
|
|
logger.info("=" * 60)
|
|
logger.info("🔄 开始自动迁移流程")
|
|
logger.info("=" * 60)
|
|
|
|
# 1. 创建迁移
|
|
if not create_migration("auto_migration"):
|
|
logger.error("❌ 自动迁移失败：无法生成迁移")
|
|
return False
|
|
|
|
# 2. 执行迁移
|
|
if not upgrade_database():
|
|
logger.error("❌ 自动迁移失败：无法执行迁移")
|
|
return False
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("✅ 自动迁移完成")
|
|
logger.info("=" * 60)
|
|
return True
|
|
|
|
|
|
def init_database():
|
|
"""初始化数据库（首次部署）"""
|
|
logger.info("=" * 60)
|
|
logger.info("🔧 初始化数据库")
|
|
logger.info("=" * 60)
|
|
|
|
# 创建初始迁移
|
|
if not create_migration("initial_migration"):
|
|
logger.warning("⚠️ 无法创建初始迁移，可能已存在")
|
|
|
|
# 执行迁移
|
|
if not upgrade_database():
|
|
logger.error("❌ 初始化失败")
|
|
return False
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("✅ 数据库初始化完成")
|
|
logger.info("=" * 60)
|
|
return True
|
|
|
|
|
|
def main():
|
|
"""主函数"""
|
|
if len(sys.argv) < 2:
|
|
print("使用方法:")
|
|
print(" python migrate.py create [message] - 创建新迁移")
|
|
print(" python migrate.py upgrade [revision] - 升级数据库（默认: head）")
|
|
print(" python migrate.py downgrade [revision] - 降级数据库（默认: -1）")
|
|
print(" python migrate.py current - 显示当前版本")
|
|
print(" python migrate.py history - 显示迁移历史")
|
|
print(" python migrate.py heads - 显示最新版本")
|
|
print(" python migrate.py stamp [revision] - 标记版本（默认: head）")
|
|
print(" python migrate.py auto - 自动迁移（生成+执行）")
|
|
print(" python migrate.py init - 初始化数据库")
|
|
sys.exit(1)
|
|
|
|
command = sys.argv[1]
|
|
|
|
if command == "create":
|
|
message = sys.argv[2] if len(sys.argv) > 2 else None
|
|
success = create_migration(message)
|
|
elif command == "upgrade":
|
|
revision = sys.argv[2] if len(sys.argv) > 2 else "head"
|
|
success = upgrade_database(revision)
|
|
elif command == "downgrade":
|
|
revision = sys.argv[2] if len(sys.argv) > 2 else "-1"
|
|
success = downgrade_database(revision)
|
|
elif command == "current":
|
|
success = show_current()
|
|
elif command == "history":
|
|
success = show_history()
|
|
elif command == "heads":
|
|
success = show_heads()
|
|
elif command == "stamp":
|
|
revision = sys.argv[2] if len(sys.argv) > 2 else "head"
|
|
success = stamp_database(revision)
|
|
elif command == "auto":
|
|
success = auto_migrate()
|
|
elif command == "init":
|
|
success = init_database()
|
|
else:
|
|
logger.error(f"❌ 未知命令: {command}")
|
|
success = False
|
|
|
|
sys.exit(0 if success else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |