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