"""数据库迁移模块。 使用方式: # 在应用启动时自动执行所有迁移 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)}