document-management/backend/dms/db_upgrade.py

"""dms-server.exe upgrade-db：生产库 alembic 迁移升级工具。

从数据库当前 alembic 版本升级到 head（0015_login_ip_restriction）。
0014/0015 的 upgrade 仅执行 CREATE TABLE / CREATE INDEX / CHECK 约束超集替换，
不删除、不修改既有业务数据；幂等，可重复执行。
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Optional


def resolve_migrations_dir() -> Path:
    """返回 migrations 目录：打包环境取 _MEIPASS 内置副本，源码环境取 backend/migrations。"""
    if getattr(sys, "frozen", False):
        base = Path(getattr(sys, "_MEIPASS", Path(sys.executable).parent))
        return base / "migrations"
    return Path(__file__).resolve().parents[1] / "migrations"


def _mask_password(url: str) -> str:
    """脱敏数据库 URL：scheme://user:***@host:port/db。"""
    if "@" not in url:
        return url
    head, _, tail = url.rpartition("@")
    if "://" in head:
        scheme, userinfo = head.split("://", 1)
        prefix = f"{scheme}://"
    else:
        scheme, userinfo = "", head
        prefix = ""
    user = userinfo.split(":", 1)[0]
    return f"{prefix}{user}:***@{tail}"


def _read_current_version(database_url: str) -> Optional[str]:
    """读取 alembic_version.version_num；无该表视为全新库返回 None。

    连接失败让异常冒出，由调用方以退出码 2 终止——绝不把连不上误判为全新库。
    """
    import sqlalchemy as sa

    engine = sa.create_engine(database_url)
    try:
        with engine.connect() as conn:
            if not sa.inspect(conn).has_table("alembic_version"):
                return None
            row = conn.execute(sa.text("SELECT version_num FROM alembic_version")).fetchone()
            return row[0] if row else None
    finally:
        engine.dispose()


def run_upgrade(assume_yes: bool = False) -> int:
    """编程式执行 alembic upgrade head，返回进程退出码（0 成功）。"""
    database_url = os.environ.get("DMS_DATABASE_URL", "").strip()
    if not database_url:
        print("错误：未配置 DMS_DATABASE_URL（请检查 dms-server.exe 同目录的 .env 文件）。")
        print("为避免误连错误数据库，升级命令拒绝在未明确指定数据库地址时执行。")
        return 2

    # 延迟导入：缺 URL 检查必须先于一切 alembic 初始化执行
    from alembic import command
    from alembic.config import Config
    from alembic.script import ScriptDirectory

    migrations_dir = resolve_migrations_dir()
    if not (migrations_dir / "env.py").exists():
        print(f"错误：未找到迁移目录 {migrations_dir}。")
        return 2

    cfg = Config()
    cfg.set_main_option("script_location", str(migrations_dir))
    cfg.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))

    script = ScriptDirectory.from_config(cfg)
    head = script.get_current_head()

    try:
        current = _read_current_version(database_url)
    except Exception as exc:  # noqa: BLE001 - 任意连接类错误统一转为退出码
        print(f"错误：无法连接数据库（{_mask_password(database_url)}）：{exc}")
        return 2

    print(f"数据库地址：{_mask_password(database_url)}")
    print(f"当前版本：{current or '（全新库，无 alembic_version 表）'}")
    print(f"目标版本：{head}")

    if current == head:
        print("数据库已是最新版本，无需升级。")
        return 0

    if current is None:
        revisions = [r.revision for r in reversed(list(script.walk_revisions("base", head)))]
    else:
        # walk_revisions 含 base 端点，切片去掉已执行的 current
        revisions = [r.revision for r in reversed(list(script.walk_revisions(current, head)))][1:]
    print("待执行迁移：")
    for rev in revisions:
        print(f"  - {rev}")
    print("以上迁移仅新增表/索引/约束，不会删除或修改既有业务数据。")

    if not assume_yes:
        answer = input("确认执行升级？(y/N)：").strip().lower()
        if answer not in ("y", "yes"):
            print("已取消。")
            return 1

    command.upgrade(cfg, "head")

    after = _read_current_version(database_url)
    print(f"升级完成，当前版本：{after}")
    return 0 if after == head else 2