xiaoxia-saas/scripts/check_migration_safety.py

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

只检查 Alembic 迁移文件的 upgrade 函数中是否包含破坏性操作：
- DROP TABLE / DROP COLUMN（高风险，直接阻断）
- 不加默认值的 NOT NULL 列添加（中风险，告警）
- 大表的 ALTER TABLE（超过10万行，中风险，告警）
- 列类型变更（可能导致数据丢失，中风险）
- 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 os
import re
import subprocess
import sys
from pathlib import Path
from typing import List, Tuple

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

# 已知可能的大表（超过10万行的表），ALTER TABLE 这些表需要告警
# 可通过环境变量 MIGRATION_LARGE_TABLES 覆盖，用逗号分隔
LARGE_TABLES = os.getenv(
    "MIGRATION_LARGE_TABLES",
    "generation_projects,generation_tasks,assets,upload_sessions,users",
).split(",")
LARGE_TABLE_THRESHOLD = int(os.getenv("MIGRATION_LARGE_TABLE_THRESHOLD", "100000"))

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

# 中风险模式：可能导致数据丢失或兼容性问题
# 注意：NOT NULL 检测需要额外判断是否有默认值，使用专用函数检测
MEDIUM_RISK_PATTERNS = [
    (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 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 extract_alter_column_calls(upgrade_content: str) -> List[str]:
    """
    从 upgrade 内容中提取所有 op.alter_column() 调用的完整参数字符串。
    处理跨行的情况。
    """
    calls = []
    pattern = r"op\.alter_column\("
    pos = 0
    while True:
        match = re.search(pattern, upgrade_content[pos:])
        if not match:
            break
        call_start = pos + match.start()
        # 从 opening 括号之后开始计数，初始深度为 1
        i = call_start + len("op.alter_column(")
        depth = 1
        while i < len(upgrade_content) and depth > 0:
            if upgrade_content[i] == "(":
                depth += 1
            elif upgrade_content[i] == ")":
                depth -= 1
            i += 1
        if depth == 0:
            # 找到了匹配的闭合括号，i 指向闭合括号之后
            calls.append(upgrade_content[call_start : i - 1])
        pos = i
    return calls


def check_not_null_without_default(file_path: Path, upgrade_content: str) -> List[str]:
    """
    检测 NOT NULL 列添加但没有默认值的情况。
    只有 nullable=False 且没有 server_default / default 时才告警。
    """
    issues = []
    alter_calls = extract_alter_column_calls(upgrade_content)

    for call in alter_calls:
        # 检查是否设置了 nullable=False
        if not re.search(r"nullable\s*=\s*False", call):
            continue

        # 检查是否有默认值（server_default 或 default）
        has_default = bool(re.search(r"server_default\s*=", call) or re.search(r"\bdefault\s*=", call))

        if not has_default:
            # 提取表名和列名
            # op.alter_column('table_name', 'column_name', ...)
            table_col_match = re.search(
                r"op\.alter_column\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]",
                call,
            )
            if table_col_match:
                table = table_col_match.group(1)
                col = table_col_match.group(2)
                issues.append(
                    f"{file_path.name}: 新增 NOT NULL 约束无默认值 " f"（表: {table}, 列: {col}）- 旧数据为空时迁移失败"
                )
            else:
                issues.append(f"{file_path.name}: 新增 NOT NULL 约束无默认值 " "- 旧数据为空时迁移失败")

    return issues


def check_large_table_alter(file_path: Path, upgrade_content: str) -> List[str]:
    """
    检测大表的 ALTER TABLE 操作。
    大表名单通过 LARGE_TABLES 配置。
    """
    issues = []

    # 检测 alter_column 在大表上的操作
    alter_calls = extract_alter_column_calls(upgrade_content)

    for call in alter_calls:
        table_col_match = re.search(
            r"op\.alter_column\(\s*['\"]([^'\"]+)['\"]",
            call,
        )
        if not table_col_match:
            continue

        table = table_col_match.group(1)
        if table in LARGE_TABLES:
            # 提取操作类型
            ops = []
            if re.search(r"nullable\s*=", call):
                ops.append("修改nullable")
            if re.search(r"type_\s*=", call):
                ops.append("列类型变更")
            if re.search(r"server_default\s*=", call):
                ops.append("修改默认值")
            if not ops:
                ops.append("ALTER COLUMN")

            issues.append(
                f"{file_path.name}: 大表 ALTER TABLE "
                f"（表: {table}, 操作: {', '.join(ops)}）"
                f" - 预计行数 > {LARGE_TABLE_THRESHOLD}，可能导致长时间锁表"
            )

    # 检测大表上的 drop_column
    drop_col_pattern = r"op\.drop_column\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]"
    for match in re.finditer(drop_col_pattern, upgrade_content):
        table = match.group(1)
        col = match.group(2)
        if table in LARGE_TABLES:
            # 这个已经算高风险了，但额外加上大表提示
            pass  # drop_column 已在 HIGH_RISK 中覆盖

    return issues


def get_new_migrations_via_diff(diff_target: str) -> List[Path]:
    """
    通过 git diff 对比目标分支/commit，找出 alembic/versions/ 下新增的迁移文件。
    只包含新增文件（A状态），不包含修改或删除的文件。
    """
    try:
        result = subprocess.run(
            [
                "git",
                "diff",
                "--name-only",
                "--diff-filter=A",
                diff_target,
                "HEAD",
                "--",
                "alembic/versions/",
            ],
            cwd=str(REPO_ROOT),
            capture_output=True,
            text=True,
            check=True,
        )
        files = [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
        return [REPO_ROOT / f for f in files]
    except subprocess.CalledProcessError as e:
        print(f"⚠️  git diff 失败（{diff_target}）：{e.stderr.strip()}")
        print(f"   降级为检查所有迁移文件")
        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}")

    # NOT NULL 无默认值检测（专用精确检测）
    not_null_issues = check_not_null_without_default(file_path, upgrade_content)
    medium_risks.extend(not_null_issues)

    # 大表 ALTER TABLE 检测
    large_table_issues = check_large_table_alter(file_path, upgrade_content)
    medium_risks.extend(large_table_issues)

    # 安全操作检测
    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()
    print(f"📋 大表名单（ALTER TABLE 会额外告警）: {', '.join(LARGE_TABLES)}")
    print(f"📊 大表阈值: > {LARGE_TABLE_THRESHOLD} 行")
    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())