AcidMap/db_migrate.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据库迁移辅助脚本

此脚本提供简单的命令行接口来执行常见的数据库迁移操作。
"""

import argparse
import os
import sys
import importlib
import traceback
from alembic.config import Config
from alembic import command

# 直接设置环境变量，确保所有地方使用相同的数据库连接信息
os.environ["DB_USER"] = "postgres"
os.environ["DB_PASSWORD"] = "123456789Qq"
os.environ["DB_HOST"] = "localhost"
os.environ["DB_PORT"] = "5432"
os.environ["DB_NAME"] = "soilgd"

def get_alembic_config():
    """
    获取 Alembic 配置
    
    Returns:
        Alembic 配置对象
    """
    # 使用相对路径获取配置文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    alembic_ini_path = os.path.join(script_dir, 'alembic.ini')
    print(f"使用配置文件: {alembic_ini_path}")
    
    # 创建配置对象
    alembic_cfg = Config(alembic_ini_path)
    return alembic_cfg


def create_migration(message, auto=True):
    """
    创建新的迁移
    
    Args:
        message: 迁移说明
        auto: 是否自动生成迁移内容
    """
    cfg = get_alembic_config()
    try:
        print(f"开始创建迁移: {message}, 自动生成: {auto}")
        if auto:
            command.revision(cfg, message=message, autogenerate=True)
        else:
            command.revision(cfg, message=message, autogenerate=False)
        print(f"成功创建迁移 '{message}'")
        return 0
    except Exception as e:
        print(f"创建迁移失败: {str(e)}")
        print("详细错误信息:")
        traceback.print_exc()
        return 1


def upgrade(target="head"):
    """
    升级数据库
    
    Args:
        target: 目标版本，默认为最新版本
    """
    cfg = get_alembic_config()
    try:
        command.upgrade(cfg, target)
        print(f"成功升级数据库到 '{target}'")
        return 0
    except Exception as e:
        print(f"升级数据库失败: {str(e)}")
        return 1


def downgrade(target="-1"):
    """
    降级数据库
    
    Args:
        target: 目标版本，默认为上一个版本
    """
    cfg = get_alembic_config()
    try:
        command.downgrade(cfg, target)
        print(f"成功降级数据库到 '{target}'")
        return 0
    except Exception as e:
        print(f"降级数据库失败: {str(e)}")
        return 1


def show_history():
    """显示迁移历史"""
    cfg = get_alembic_config()
    try:
        command.history(cfg)
        return 0
    except Exception as e:
        print(f"显示历史失败: {str(e)}")
        return 1


def show_current():
    """显示当前版本"""
    cfg = get_alembic_config()
    try:
        command.current(cfg)
        return 0
    except Exception as e:
        print(f"显示当前版本失败: {str(e)}")
        return 1


def stamp_revision(target="head"):
    """
    将数据库版本标记为指定版本，而不运行任何迁移
    
    Args:
        target: 目标版本，默认为最新版本
    """
    cfg = get_alembic_config()
    try:
        command.stamp(cfg, target)
        print(f"成功将数据库版本标记为 '{target}'")
        return 0
    except Exception as e:
        print(f"标记数据库版本失败: {str(e)}")
        return 1


def execute_sql(sql_statement):
    """
    直接执行 SQL 语句
    
    Args:
        sql_statement: 要执行的 SQL 语句
    """
    try:
        from app.database import execute_sql
        execute_sql(sql_statement)
        print(f"成功执行 SQL: {sql_statement}")
        return 0
    except Exception as e:
        print(f"执行 SQL 失败: {str(e)}")
        return 1


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='数据库迁移辅助工具')
    subparsers = parser.add_subparsers(dest='command', help='命令')
    
    # 检查 alembic 是否已安装
    try:
        importlib.import_module('alembic')
    except ImportError:
        print("错误: 未找到 alembic 模块，请先安装: pip install alembic")
        return 1
    
    # 创建迁移
    create_parser = subparsers.add_parser('create', help='创建新的迁移')
    create_parser.add_argument('message', help='迁移说明')
    create_parser.add_argument('--manual', action='store_true', help='手动创建迁移（不自动生成）')
    
    # 升级数据库
    upgrade_parser = subparsers.add_parser('upgrade', help='升级数据库')
    upgrade_parser.add_argument('target', nargs='?', default='head', help='目标版本，默认为最新版本')
    
    # 降级数据库
    downgrade_parser = subparsers.add_parser('downgrade', help='降级数据库')
    downgrade_parser.add_argument('target', nargs='?', default='-1', help='目标版本，默认为上一个版本')
    
    # 标记数据库版本
    stamp_parser = subparsers.add_parser('stamp', help='标记数据库版本，不执行迁移')
    stamp_parser.add_argument('target', nargs='?', default='head', help='目标版本，默认为最新版本')
    
    # 执行 SQL
    sql_parser = subparsers.add_parser('sql', help='执行 SQL 语句')
    sql_parser.add_argument('statement', help='要执行的 SQL 语句')
    
    # 显示历史
    subparsers.add_parser('history', help='显示迁移历史')
    
    # 显示当前版本
    subparsers.add_parser('current', help='显示当前版本')
    
    args = parser.parse_args()
    
    if args.command == 'create':
        return create_migration(args.message, not args.manual)
    elif args.command == 'upgrade':
        return upgrade(args.target)
    elif args.command == 'downgrade':
        return downgrade(args.target)
    elif args.command == 'stamp':
        return stamp_revision(args.target)
    elif args.command == 'sql':
        return execute_sql(args.statement)
    elif args.command == 'history':
        return show_history()
    elif args.command == 'current':
        return show_current()
    else:
        parser.print_help()
        return 0


if __name__ == '__main__':
    sys.exit(main())