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

    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()


def _acquire_advisory_lock():
    """PostgreSQL advisory lock，防止多 Worker 并发执行迁移。

    返回 True 表示获得锁，False 表示另一个进程正在迁移。
    """
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        result = db.session.execute(text("SELECT pg_try_advisory_lock(20260727)"))
        locked = result.scalar()
        return bool(locked)
    except Exception:
        # SQLite 或其他数据库不支持 advisory lock，用文件锁替代
        import fcntl
        lock_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".migration.lock")
        try:
            lock_fd = open(lock_path, "w")
            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
            return True
        except (IOError, OSError):
            return False


def _release_advisory_lock():
    """释放 advisory lock。"""
    from insurance.db.compat import db
    from sqlalchemy import text

    try:
        db.session.execute(text("SELECT pg_advisory_unlock(20260727)"))
    except Exception:
        pass


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

    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()


def run_migrations():
    """自动执行所有迁移脚本（带幂等保护、事务回滚和并发锁）。"""
    from flask import current_app
    from insurance.db.compat import db

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

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

    # 获取 advisory lock，防止多 Worker 并发迁移
    if not _acquire_advisory_lock():
        logger.info("另一个进程正在执行迁移，跳过")
        return

    try:
        # 确保迁移历史表存在
        _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])

        # 按顺序执行未执行的迁移
        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)
                    logger.info(f"迁移 {migration_name} 完成")
            except Exception as e:
                # 回滚残留事务
                try:
                    db.session.rollback()
                except Exception:
                    pass
                logger.error(f"迁移 {migration_name} 失败: {e}")
                # 失败后停止，不带病继续
                raise RuntimeError(f"迁移 {migration_name} 失败，中止启动: {e}") from e

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


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)}