#!/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()