baodan/api/insurance/db/__init__.py

"""数据库迁移模块。

使用方式：
    # 在应用启动时自动执行所有迁移
    from insurance.db import run_migrations
    run_migrations()

    # 或手动执行单个迁移
    from insurance.db.migrate_001 import migrate
    migrate()
"""
import importlib
import logging
import os
from datetime import datetime

logger = logging.getLogger(__name__)


def _ensure_migration_table():
    """确保迁移历史表存在。"""
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        # 创建迁移历史表（如果不存在）
        db.session.execute(text("""
            CREATE TABLE IF NOT EXISTS db_migration_history (
                id SERIAL PRIMARY KEY,
                migration_name VARCHAR(255) UNIQUE NOT NULL,
                executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """))
        db.session.commit()
    except Exception as e:
        logger.warning(f"创建迁移历史表失败: {e}")


def _get_executed_migrations():
    """获取已执行的迁移列表。"""
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        result = db.session.execute(text(
            "SELECT migration_name FROM db_migration_history ORDER BY executed_at"
        ))
        return {row[0] for row in result}
    except Exception:
        return set()


def _record_migration(migration_name: str):
    """记录迁移已执行。"""
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        db.session.execute(text(
            "INSERT INTO db_migration_history (migration_name, executed_at) VALUES (:name, :now)"
        ), {"name": migration_name, "now": datetime.now()})
        db.session.commit()
    except Exception as e:
        logger.warning(f"记录迁移 {migration_name} 失败: {e}")
        try:
            db.session.rollback()
        except Exception:
            pass


def run_migrations():
    """自动执行所有迁移脚本（带幂等保护）。"""
    from flask import current_app

    # 检查是否启用自动迁移
    if not current_app.config.get("MIGRATION_ENABLED", False):
        return

    logger.info("开始执行数据库迁移...")

    # 确保迁移历史表存在
    _ensure_migration_table()

    # 获取已执行的迁移
    executed = _get_executed_migrations()

    # 获取当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))

    # 扫描所有迁移脚本（migrate_*.py）
    migrations = []
    for filename in sorted(os.listdir(current_dir)):
        if filename.startswith("migrate_") and filename.endswith(".py"):
            migrations.append(filename[:-3])  # 去掉 .py 后缀

    # 按顺序执行未执行的迁移
    for migration_name in migrations:
        if migration_name in executed:
            logger.debug(f"跳过已执行的迁移: {migration_name}")
            continue
        try:
            module = importlib.import_module(f"insurance.db.{migration_name}")
            if hasattr(module, "migrate"):
                logger.info(f"执行迁移: {migration_name}")
                module.migrate()
                # 记录迁移已执行
                _record_migration(migration_name)
        except Exception as e:
            logger.warning(f"迁移 {migration_name} 失败: {e}")
            # 继续执行下一个迁移，不中断

    logger.info("迁移执行完成")


def get_migration_status():
    """获取迁移执行状态。"""
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        # 检查迁移历史表是否存在
        result = db.session.execute(text("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables
                WHERE table_name = 'db_migration_history'
            )
        """))

        if not result.fetchone()[0]:
            return {"status": "no_history_table", "migrations": []}

        # 获取已执行的迁移
        result = db.session.execute(text("""
            SELECT migration_name, executed_at
            FROM db_migration_history
            ORDER BY executed_at DESC
        """))

        migrations = [dict(row) for row in result]
        return {"status": "ok", "migrations": migrations}

    except Exception as e:
        return {"status": "error", "message": str(e)}