xiaoxia-saas/scripts/ci/check_migration_chain.py

#!/usr/bin/env python3
"""
检查 Alembic migration 编号连续性。

扫描 alembic/versions/ 下所有 migration 文件，提取 revision 和 down_revision，
验证整条链是否完整——每个 down_revision（除了 baseline 的 None）都必须对应一个存在的 revision。

支持两种格式:
    revision: str = "001"          # 旧格式（带类型注解）
    revision = "038_error_retry"   # 新格式（带描述后缀）

匹配策略：提取 revision 名称的数字前缀（如 "001"、"038"）作为唯一标识进行匹配，
兼容纯数字编号和"数字_描述"两种命名风格。

用法:
    python3 scripts/ci/check_migration_chain.py [alembic_versions_dir]

默认目录: alembic/versions/

退出码:
    0 - 链完整
    1 - 有断链或其他错误
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

# 匹配 revision / down_revision，支持带类型注解和不带类型注解两种格式
# revision: str = "xxx" 或 revision = "xxx"
REV_PATTERN = re.compile(
    r'^\s*revision\s*(?::\s*str\s*)?=\s*["\']([^"\']+)["\']',
    re.MULTILINE,
)
DOWN_PATTERN = re.compile(
    r'^\s*down_revision\s*(?::\s*(?:Union\[str,\s*None\]|str\s*\|\s*None|None|str)\s*)?=\s*(["\']([^"\']+)["\']|None)',
    re.MULTILINE,
)

# 提取 revision 名称的数字前缀，如 "001" 或 "038_error_retry" → "038"
NUM_PREFIX_PATTERN = re.compile(r"^(\d+)")


def num_prefix(name: str) -> str:
    """提取 revision 名称的数字前缀。"""
    m = NUM_PREFIX_PATTERN.match(name)
    return m.group(1) if m else name


def extract_migration_info(filepath: Path) -> tuple[str, str | None]:
    """从 migration 文件中提取 revision 和 down_revision（返回完整名称）。"""
    content = filepath.read_text(encoding="utf-8")

    rev_match = REV_PATTERN.search(content)
    down_match = DOWN_PATTERN.search(content)

    if not rev_match:
        raise ValueError(f"{filepath.name}: 未找到 revision 定义")

    revision = rev_match.group(1)

    if not down_match:
        raise ValueError(f"{filepath.name}: 未找到 down_revision 定义")

    # down_match group(2) 是引号内的值，如果是 None 则 group(2) 为 None
    down_revision = down_match.group(2)

    return revision, down_revision


def check_chain(versions_dir: Path) -> list[str]:
    """检查 migration 链是否完整，返回错误列表。"""
    errors: list[str] = []

    if not versions_dir.is_dir():
        return [f"目录不存在: {versions_dir}"]

    py_files = sorted(versions_dir.glob("*.py"))
    if not py_files:
        return [f"目录下没有 migration 文件: {versions_dir}"]

    # 收集所有 revision（用数字前缀做唯一标识）
    revisions_by_num: dict[str, str] = {}  # 数字前缀 -> 完整 revision 名
    revision_files: dict[str, str] = {}  # 数字前缀 -> 文件名
    down_revisions: list[tuple[str, str | None]] = []  # (文件名, down_revision 数字前缀或None)

    for f in py_files:
        if f.name.startswith("__"):
            continue
        try:
            rev, down = extract_migration_info(f)
        except ValueError as e:
            errors.append(str(e))
            continue

        rev_num = num_prefix(rev)

        if rev_num in revisions_by_num:
            errors.append(
                f"编号重复: 编号 {rev_num} 同时出现在 "
                f"{f.name} (revision={rev}) 和 {revision_files[rev_num]} (revision={revisions_by_num[rev_num]})"
            )
        else:
            revisions_by_num[rev_num] = rev
            revision_files[rev_num] = f.name

        down_num = num_prefix(down) if down else None
        down_revisions.append((f.name, down_num))

    if errors:
        return errors

    # 检查每个 down_revision 是否存在
    baselines = 0
    for filename, down_num in down_revisions:
        if down_num is None:
            baselines += 1
            continue

        if down_num not in revisions_by_num:
            errors.append(
                f"断链: {filename} 的 down_revision 指向编号 '{down_num}'，但没有任何 migration 的 revision 是这个编号"
            )

    if baselines == 0:
        errors.append("没有找到 baseline migration（down_revision = None 的文件）")
    elif baselines > 1:
        errors.append(f"发现 {baselines} 个 baseline migration，通常只能有 1 个")

    # 额外检查：数字编号是否连续（只对能提取出数字的）
    if revisions_by_num and not errors:
        nums = sorted(int(n) for n in revisions_by_num if n.isdigit())
        if nums:
            expected = list(range(nums[0], nums[-1] + 1))
            missing = [n for n in expected if n not in nums]
            if missing:
                missing_str = ", ".join(f"{n:03d}" for n in missing)
                errors.append(f"编号不连续: 缺少编号 {missing_str}")

    return errors


def main() -> int:
    if len(sys.argv) > 1:
        versions_dir = Path(sys.argv[1])
    else:
        versions_dir = Path("alembic/versions")

    print(f"检查 migration 编号连续性: {versions_dir}")
    print()

    errors = check_chain(versions_dir)

    py_files = [f for f in versions_dir.glob("*.py") if not f.name.startswith("__")]

    if errors:
        print(f"❌ Migration 链有问题（共 {len(py_files)} 个文件，{len(errors)} 个错误）:")
        for e in errors:
            print(f"  - {e}")
        print()
        print("请修复后再提交。常见原因：")
        print("  1. 新 migration 的 down_revision 编号写错了")
        print("  2. 多个 PR 同时加 migration，编号冲突")
        print("  3. 合并代码时漏了某个 migration 文件")
        return 1

    print(f"✅ Migration 链完整，共 {len(py_files)} 个版本")
    return 0


if __name__ == "__main__":
    sys.exit(main())