xiaoxia-saas/scripts/check_migration_safety.py

#!/usr/bin/env python3
"""
数据库迁移破坏性变更安全检查

只检查 Alembic 迁移文件的 upgrade 函数中是否包含破坏性操作：
- DROP TABLE
- ALTER TABLE ... DROP COLUMN
- 列类型变更（可能导致数据丢失）
- NOT NULL 约束新增（无默认值时）
- RENAME TABLE / RENAME COLUMN

忽略 downgrade 函数中的操作（那是回滚逻辑，正常的）。

使用方式：
  # 检查所有迁移（不推荐，会扫历史已执行的迁移）
  python3 scripts/check_migration_safety.py

  # 只检查与目标分支相比新增的迁移（推荐用于CI）
  python3 scripts/check_migration_safety.py --diff-against origin/main

  # 只检查指定版本之后的迁移
  python3 scripts/check_migration_safety.py --since 030_xxx

退出码：
  0 - 安全 / 只有非破坏性变更
  1 - 检测到高风险破坏性变更
  2 - 检测到中风险变更，需人工确认
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import List, Tuple

REPO_ROOT = Path(__file__).resolve().parents[1]
ALEMBIC_VERSIONS_DIR = REPO_ROOT / "alembic" / "versions"

# 高风险模式：直接导致数据丢失（只在 upgrade 中检查）
HIGH_RISK_PATTERNS = [
    (r"\bop\.drop_table\(", "op.drop_table() - 删除表，数据永久丢失"),
    (r"\bop\.drop_column\(", "op.drop_column() - 删除列，数据永久丢失"),
]

# 中风险模式：可能导致数据丢失或兼容性问题
MEDIUM_RISK_PATTERNS = [
    (
        r"op\.alter_column\([^)]*nullable\s*=\s*False",
        "新增 NOT NULL 约束 - 旧数据可能为空导致迁移失败",
    ),
    (r"op\.alter_column\([^)]*type_\s*=", "列类型变更 - 可能导致数据截断或转换失败"),
    (
        r"\bop\.rename_table\(",
        "op.rename_table() - 重命名表，可能导致依赖该表的代码报错",
    ),
    (
        r"\bop\.rename_column\(",
        "op.rename_column() - 重命名列，可能导致依赖该列的代码报错",
    ),
    (r"\bop\.drop_index\(", "op.drop_index() - 删除索引，可能影响查询性能"),
    (r"\bop\.drop_constraint\(", "op.drop_constraint() - 删除约束，可能影响数据完整性"),
]

# 安全模式：这些是安全的新增操作
SAFE_PATTERNS = [
    (r"\bop\.create_table\(", "新建表"),
    (r"\bop\.add_column\(", "新增列"),
    (r"\bop\.create_index\(", "新建索引"),
    (r"\bop\.create_unique_constraint\(", "新建唯一约束"),
    (r"\bop\.create_foreign_key\(", "新建外键约束"),
]


def _get_env(*names: str, default: str = "") -> str:
    """按优先级尝试多个环境变量名，返回第一个非空值。"""
    for name in names:
        val = os.environ.get(name, "")
        if val:
            return val
    return default


def extract_upgrade_content(content: str) -> str:
    """
    从迁移文件中提取 upgrade 函数的内容。
    只检查 upgrade 中的操作，忽略 downgrade。
    """
    upgrade_match = re.search(r"def upgrade\b[^:]*:", content)
    if not upgrade_match:
        return ""

    upgrade_start = upgrade_match.end()

    # 找到下一个顶层 def（通常是 def downgrade）作为结束位置
    rest = content[upgrade_start:]
    downgrade_match = re.search(r"\n\ndef\s+\w+\b", rest)
    if downgrade_match:
        upgrade_end = upgrade_start + downgrade_match.start()
    else:
        upgrade_end = len(content)

    return content[upgrade_start:upgrade_end]


def _api_get_with_retry(url: str, token: str, max_retries: int = 3) -> dict | list:
    """
    带重试的 API 调用。
    指数退避：1s, 2s, 4s
    """
    last_error = None
    for attempt in range(max_retries):
        try:
            req = urllib.request.Request(url, headers={"Authorization": f"token {token}"})
            with urllib.request.urlopen(req, timeout=15) as resp:
                return json.loads(resp.read().decode())
        except urllib.error.HTTPError as e:
            # 404 说明目录不存在或分支不存在，直接抛
            if e.code == 404:
                raise
            last_error = e
            if attempt < max_retries - 1:
                wait = 2**attempt
                print(f"   （API 请求失败，{wait}s 后重试 {attempt + 1}/{max_retries}：{e}）")
                time.sleep(wait)
        except Exception as e:
            last_error = e
            if attempt < max_retries - 1:
                wait = 2**attempt
                print(f"   （API 请求失败，{wait}s 后重试 {attempt + 1}/{max_retries}：{e}）")
                time.sleep(wait)
    raise last_error  # type: ignore[misc]


def get_new_migrations_via_api(diff_target: str) -> List[Path] | None:
    """
    通过 Gitea/GitHub Contents API 对比目标分支，找出 alembic/versions/ 下新增的迁移文件。
    返回 None 表示 API 方式不可用，调用方应尝试其他方式。
    """
    # 同时支持 Gitea 和 GitHub 的环境变量命名
    api_url = _get_env("GITEA_API_URL", "GITHUB_API_URL", "CI_API_V4_URL")
    repo = _get_env("GITEA_REPOSITORY", "GITHUB_REPOSITORY", "CI_PROJECT_PATH")
    token = _get_env("GITEA_TOKEN", "GITHUB_TOKEN", "CI_JOB_TOKEN")
    branch = diff_target.replace("origin/", "")

    if not api_url or not repo or not token:
        print(
            f"   （API 环境变量不完整：api_url={'✓' if api_url else '✗'} repo={'✓' if repo else '✗'} token={'✓' if token else '✗'}）"
        )
        return None

    try:
        url = f"{api_url}/repos/{repo}/contents/alembic/versions?ref={branch}"
        data = _api_get_with_retry(url, token)

        if isinstance(data, dict):
            # Gitea 目录不存在时返回 404，不会到这里；如果返回 dict 可能是错误信息
            print(f"   （API 返回异常：{str(data)[:100]}）")
            return None

        remote_files = {item["name"] for item in data if item["name"].endswith(".py")}
        local_files = {f.name for f in ALEMBIC_VERSIONS_DIR.glob("*.py")}
        new_file_names = sorted(local_files - remote_files)

        if new_file_names:
            result = [ALEMBIC_VERSIONS_DIR / f for f in new_file_names]
            print(f"   （API 对比 {branch} 分支，发现 {len(result)} 个新增迁移）")
            return result
        else:
            print(f"   （API 对比 {branch} 分支，无新增迁移）")
            return []
    except Exception as e:
        print(f"   （API 获取迁移列表失败：{e}）")
        return None


def get_new_migrations_via_git(diff_target: str) -> List[Path] | None:
    """
    Fallback：通过本地 git diff 找出新增的迁移文件。
    CI 环境中 git 可用时作为 API 失败后的兜底方案。
    """
    try:
        # 确保目标分支存在
        subprocess.run(
            ["git", "fetch", "origin", diff_target.replace("origin/", ""), "--depth=50"],
            capture_output=True,
            cwd=str(REPO_ROOT),
            timeout=30,
        )

        # 优先使用三点diff（找合并基线），失败时回退到两点diff（兼容tar.gz checkout + git init的CI环境）
        diff_args = ["git", "diff", "--name-only", "--diff-filter=A", f"{diff_target}...HEAD"]
        result = subprocess.run(
            diff_args,
            capture_output=True,
            text=True,
            cwd=str(REPO_ROOT),
            timeout=10,
        )
        if result.returncode != 0:
            # fallback: 两点diff（无需共同祖先）
            diff_args_2 = ["git", "diff", "--name-only", "--diff-filter=A", diff_target, "HEAD"]
            result = subprocess.run(
                diff_args_2,
                capture_output=True,
                text=True,
                cwd=str(REPO_ROOT),
                timeout=10,
            )
            if result.returncode != 0:
                print(f"   （git diff 失败：{result.stderr.strip()}）")
                return None

        new_migrations = []
        for line in result.stdout.strip().split("\n"):
            line = line.strip()
            if line.startswith("alembic/versions/") and line.endswith(".py"):
                new_migrations.append(REPO_ROOT / line)

        new_migrations.sort()
        print(f"   （git diff 对比 {diff_target}，发现 {len(new_migrations)} 个新增迁移）")
        return new_migrations
    except Exception as e:
        print(f"   （git diff 方式失败：{e}）")
        return None


def get_new_migrations_via_diff(diff_target: str) -> List[Path]:
    """
    找出相对目标分支新增的迁移文件，按优先级尝试多种方式：
    1. Gitea/GitHub Contents API（最可靠，不受本地 checkout 深度影响）
    2. git diff（API 失败时的兜底）
    3. 全量扫描（以上都失败时的最后兜底，会输出警告）
    """
    print("🔍 尝试通过 API 获取新增迁移列表...")
    result = get_new_migrations_via_api(diff_target)
    if result is not None:
        return result

    print("🔍 API 不可用，尝试 git diff 方式...")
    result = get_new_migrations_via_git(diff_target)
    if result is not None:
        return result

    print("⚠️  所有增量方式均失败，降级为检查所有迁移文件")
    print("   这可能导致历史迁移中的破坏性操作被误报")
    print("   建议检查 CI 环境变量配置（GITHUB_API_URL / GITHUB_REPOSITORY / GITHUB_TOKEN）")
    return sorted(ALEMBIC_VERSIONS_DIR.glob("*.py"))


def find_new_migrations(since_revision: str | None = None, diff_against: str | None = None) -> List[Path]:
    """
    找出需要检查的迁移文件。
    优先级：diff_against > since_revision > 全部
    """
    if diff_against:
        return get_new_migrations_via_diff(diff_against)

    all_migrations = sorted(ALEMBIC_VERSIONS_DIR.glob("*.py"))
    if not since_revision:
        return all_migrations

    result = []
    found = False
    for m in all_migrations:
        if since_revision in m.name or since_revision in m.stem:
            found = True
            continue
        if found:
            result.append(m)

    return result if found else all_migrations


def analyze_migration(file_path: Path) -> Tuple[List[str], List[str], List[str]]:
    """分析单个迁移文件 upgrade 部分的风险等级"""
    content = file_path.read_text()
    upgrade_content = extract_upgrade_content(content)

    if not upgrade_content:
        return [], [], [f"{file_path.name}: 未找到 upgrade 函数"]

    high_risks = []
    medium_risks = []
    safes = []

    for pattern, desc in HIGH_RISK_PATTERNS:
        if re.search(pattern, upgrade_content):
            high_risks.append(f"{file_path.name}: {desc}")

    for pattern, desc in MEDIUM_RISK_PATTERNS:
        if re.search(pattern, upgrade_content):
            medium_risks.append(f"{file_path.name}: {desc}")

    for pattern, desc in SAFE_PATTERNS:
        if re.search(pattern, upgrade_content):
            safes.append(f"{file_path.name}: {desc}")

    return high_risks, medium_risks, safes


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--since",
        default=os.getenv("MIGRATION_SINCE_REVISION"),
        help="只检查指定版本之后的迁移（如：030_xxx），不传则检查所有迁移",
    )
    parser.add_argument(
        "--diff-against",
        default=os.getenv("MIGRATION_DIFF_AGAINST"),
        help="对比指定分支/commit，只检查新增的迁移文件（推荐用于CI，如 origin/main）",
    )
    parser.add_argument(
        "--warn-only",
        action="store_true",
        help="只警告不失败（用于非强制门禁场景）",
    )
    parser.add_argument(
        "--allow-medium-risk",
        action="store_true",
        help="允许中风险变更（只拦截高风险）",
    )
    args = parser.parse_args()

    migrations = find_new_migrations(args.since, args.diff_against)

    if not migrations:
        print("✅ 未找到需要检查的新增迁移文件，跳过")
        return 0

    print(f"🔍 正在检查 {len(migrations)} 个迁移文件的 upgrade 操作...")
    if args.diff_against:
        print(f"   （对比基准：{args.diff_against}，仅检查新增迁移）")
    print()

    all_high = []
    all_medium = []
    all_safe = []

    for m in migrations:
        high, medium, safe = analyze_migration(m)
        all_high.extend(high)
        all_medium.extend(medium)
        all_safe.extend(safe)

    if all_safe:
        print("✅ 安全变更：")
        for s in all_safe:
            print(f"   - {s}")
        print()

    if all_medium:
        print("⚠️  中风险变更（需人工确认）：")
        for m_item in all_medium:
            print(f"   - {m_item}")
        print()

    if all_high:
        print("❌ 高风险破坏性变更（禁止自动部署）：")
        for h in all_high:
            print(f"   - {h}")
        print()

    print("=" * 60)
    print(f"检查结果：{len(all_safe)} 项安全 / {len(all_medium)} 项中风险 / {len(all_high)} 项高风险")
    print()

    if all_high:
        print("❌ 检测到高风险破坏性变更，CI 检查失败！")
        print("   如果确认这是预期操作，请在 MR/PR 中说明原因并获得审批。")
        if args.warn_only:
            return 0
        return 1

    if all_medium and not args.allow_medium_risk:
        print("⚠️  检测到中风险变更，请人工确认后再部署。")
        if args.warn_only:
            return 0
        print("（如需仅拦截高风险，可使用 --allow-medium-risk 参数）")
        return 2

    print("✅ 未检测到破坏性变更")
    return 0


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