MoviePilot/app/startup/database_initializer.py

from collections.abc import Callable
from configparser import ConfigParser as _ConfigParser
import traceback

from alembic.command import upgrade
from alembic.config import Config
from alembic.migration import MigrationContext
from alembic.script import ScriptDirectory
from alembic.util import CommandError
from sqlalchemy import inspect
from sqlalchemy.engine import Engine

from app.runtime.config import settings
from app.db.base import Base
from app.db.engine import get_engine
from app.db.models import load_all_models
from app.runtime.log import logger
from app.startup.database import build_database_governance


def _build_alembic_config(engine: Engine | None = None) -> Config:
    """构造与应用活动数据库一致的 Alembic 配置。"""
    engine = engine or get_engine()
    alembic_cfg = Config()
    alembic_cfg.file_config = _ConfigParser(interpolation=None)
    alembic_cfg.set_main_option(
        'script_location',
        str(settings.ROOT_PATH / 'database'),
    )
    alembic_cfg.set_main_option(
        'sqlalchemy.url',
        engine.url.render_as_string(hide_password=False),
    )
    return alembic_cfg


def _migration_state(
        engine: Engine,
        alembic_cfg: Config,
) -> tuple[bool, tuple[str, ...], tuple[str, ...]]:
    """读取数据库迁移状态，并在结构写入前校验版本链。"""
    script = ScriptDirectory.from_config(alembic_cfg)
    target_heads = tuple(script.get_heads())
    with engine.connect() as connection:
        table_names = set(inspect(connection).get_table_names())
        current_heads = tuple(
            MigrationContext.configure(connection).get_current_heads()
        )
    has_existing_database = bool(table_names - {'alembic_version'})
    _validate_migration_lineage(script, current_heads, target_heads)
    return has_existing_database, current_heads, target_heads


def _validate_migration_lineage(
        script: ScriptDirectory,
        current_heads: tuple[str, ...],
        target_heads: tuple[str, ...],
) -> None:
    """拒绝无法沿当前迁移链安全升级的数据库版本。"""
    if len(target_heads) != 1:
        raise RuntimeError(
            f"数据库迁移脚本必须只有一个 head，当前为 {target_heads}"
        )
    if len(current_heads) > 1:
        raise RuntimeError(
            f"数据库存在多个 current revision，无法自动迁移：{current_heads}"
        )
    if not current_heads:
        return

    current = current_heads[0]
    target = target_heads[0]
    try:
        script.get_revision(current)
    except CommandError as error:
        raise RuntimeError(
            f"当前 MoviePilot 无法识别数据库 revision：{current}"
        ) from error
    if current == target:
        return

    ancestors = {
        revision.revision
        for revision in script.walk_revisions(base='base', head=target)
    }
    if current not in ancestors:
        raise RuntimeError(
            f"数据库 revision {current} 不是当前 head {target} 的可升级祖先"
        )


def prepare_database(*, before_alembic: Callable[[], None] | None = None) -> None:
    """在建表或迁移前完成版本校验及可选备份。"""
    engine = get_engine()
    alembic_cfg = _build_alembic_config(engine)
    has_existing_database, current_heads, target_heads = _migration_state(
        engine,
        alembic_cfg,
    )
    requires_migration = (
        has_existing_database
        and set(current_heads) != set(target_heads)
    )
    if (
        requires_migration
        and settings.DB_BACKUP_ENABLE
        and settings.DB_BACKUP_ON_UPGRADE
    ):
        current_version = current_heads[0] if current_heads else "未标记"
        target_version = target_heads[0]
        logger.info(
            f"数据库需要从版本 {current_version} 升级到 {target_version}，"
            "正在创建迁移前备份"
        )
        build_database_governance().create_backup()

    init_db()
    if before_alembic:
        # 首次初始化需要先建立用户表，再把管理员密码交给 Alembic 基础迁移消费。
        before_alembic()
    update_db(alembic_cfg)


def verify_database_revision() -> None:
    """确认活动数据库已位于当前唯一 Alembic head，否则阻止 readiness。"""
    engine = get_engine()
    alembic_cfg = _build_alembic_config(engine)
    _, current_heads, target_heads = _migration_state(engine, alembic_cfg)
    if set(current_heads) != set(target_heads):
        raise RuntimeError(
            "数据库迁移完成后 revision 仍未到达当前 head："
            f"current={current_heads}, target={target_heads}"
        )


def init_db():
    """
    初始化数据库
    """
    # 确保所有模型都已注册到 Base.metadata 中
    load_all_models()

    # 全量建表
    Base.metadata.create_all(bind=get_engine())


def update_db(alembic_cfg: Config | None = None):
    """
    更新数据库
    """
    try:
        alembic_cfg = alembic_cfg or _build_alembic_config()
        upgrade(alembic_cfg, 'head')
    except Exception as error:
        logger.error(
            f"数据库更新失败：{error}\n{traceback.format_exc()}"
        )
        raise