MoviePilot/app/startup/initializers/database.py

"""
初始化数据库结构并执行启动期迁移。
"""

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

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.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.runtime.settings import get_runtime_setting
from app.startup.composition.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(get_runtime_setting('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)
    )
    migrate_before_create = requires_migration and bool(current_heads)
    if (
        requires_migration
        and get_runtime_setting('DB_BACKUP_ENABLE')
        and get_runtime_setting('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()

    if migrate_before_create:
        # 已标记的旧结构必须先沿 Alembic 链升级；当前元数据可能包含依赖新列的外键，
        # 提前 create_all 会在 PostgreSQL 上因父表仍是旧结构而失败。
        update_db(alembic_cfg)

    init_db()
    if before_alembic:
        # 首次初始化需要先建立用户表，再把管理员密码交给 Alembic 基础迁移消费。
        before_alembic()
    if not migrate_before_create:
        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() -> None:
    """
    初始化数据库
    """
    # 确保所有模型都已注册到 Base.metadata 中
    load_all_models()

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


def update_db(alembic_cfg: Config | None = 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