60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import os
|
|
import sys
|
|
import argparse
|
|
from alembic.config import Config
|
|
from alembic import command
|
|
|
|
def create_migration(message=None):
|
|
"""
|
|
创建数据库迁移
|
|
|
|
Args:
|
|
message: 迁移消息，描述此次迁移的内容
|
|
"""
|
|
# 获取Alembic配置
|
|
alembic_cfg = Config("alembic.ini")
|
|
|
|
# 创建迁移
|
|
if message:
|
|
command.revision(alembic_cfg, message=message, autogenerate=True)
|
|
else:
|
|
command.revision(alembic_cfg, message="自动生成的迁移", autogenerate=True)
|
|
|
|
def upgrade_database():
|
|
"""更新数据库到最新版本"""
|
|
alembic_cfg = Config("alembic.ini")
|
|
command.upgrade(alembic_cfg, "head")
|
|
|
|
def downgrade_database(revision):
|
|
"""回滚数据库到指定版本"""
|
|
alembic_cfg = Config("alembic.ini")
|
|
command.downgrade(alembic_cfg, revision)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="数据库迁移工具")
|
|
subparsers = parser.add_subparsers(dest="command", help="命令")
|
|
|
|
# 创建迁移命令
|
|
create_parser = subparsers.add_parser("create", help="创建数据库迁移")
|
|
create_parser.add_argument("-m", "--message", help="迁移描述消息")
|
|
|
|
# 更新数据库命令
|
|
upgrade_parser = subparsers.add_parser("upgrade", help="更新数据库到最新版本")
|
|
|
|
# 回滚数据库命令
|
|
downgrade_parser = subparsers.add_parser("downgrade", help="回滚数据库到指定版本")
|
|
downgrade_parser.add_argument("revision", help="目标版本，例如：-1表示回滚一个版本")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == "create":
|
|
create_migration(args.message)
|
|
print("迁移文件已创建，请查看 alembic/versions 目录")
|
|
elif args.command == "upgrade":
|
|
upgrade_database()
|
|
print("数据库已更新到最新版本")
|
|
elif args.command == "downgrade":
|
|
downgrade_database(args.revision)
|
|
print(f"数据库已回滚到版本: {args.revision}")
|
|
else:
|
|
parser.print_help() |