ai-exchange/scripts/run_migration_static_checks.py

from pathlib import Path
import json
import os
import re
import sys

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.run_autonomous_cycle import REQUIRED_PIPELINE_SCHEMA

MIGRATIONS_DIR = ROOT / "migrations"
REQUIRED_SEED_VALUES = {
    "venues.code": ("OKX",),
    "instruments.symbol": ("BTC-USDT",),
}


def run_migration_static_checks(migrations_dir: Path = MIGRATIONS_DIR, required_schema=None, required_seed_values=None) -> dict:
    using_default_schema = required_schema is None
    required_schema = required_schema or REQUIRED_PIPELINE_SCHEMA
    if required_seed_values is None:
        required_seed_values = REQUIRED_SEED_VALUES if using_default_schema else {}
    files = iter_migration_files(migrations_dir)
    errors = []
    errors.extend(validate_migration_sequence(files))
    created_schema = extract_created_schema(files)

    for table, required_columns in required_schema.items():
        if table not in created_schema:
            errors.append(f"missing_table:{table}")
            continue
        missing_columns = [column for column in required_columns if column not in created_schema[table]]
        for column in missing_columns:
            errors.append(f"missing_column:{table}.{column}")
    errors.extend(validate_required_seed_values(files=files, required_seed_values=required_seed_values))

    return {
        "passed": len(errors) == 0,
        "errors": errors,
        "migration_count": len(files),
        "tables_checked": sorted(required_schema),
        "seed_values_checked": {
            key: list(values)
            for key, values in required_seed_values.items()
        },
    }


def iter_migration_files(migrations_dir: Path = MIGRATIONS_DIR) -> list[Path]:
    return sorted(path for path in migrations_dir.glob("*.sql") if path.is_file())


def validate_migration_sequence(files: list[Path]) -> list[str]:
    errors = []
    numbers = []
    for path in files:
        match = re.match(r"^(\d{3})_", path.name)
        if not match:
            errors.append(f"invalid_migration_name:{path.name}")
            continue
        numbers.append(int(match.group(1)))
    expected = list(range(1, len(numbers) + 1))
    if numbers != expected:
        errors.append(f"non_contiguous_migration_sequence:{numbers}")
    return errors


def extract_created_schema(files: list[Path]) -> dict[str, set[str]]:
    schema: dict[str, set[str]] = {}
    for path in files:
        sql = path.read_text(encoding="utf-8")
        schema.update(extract_created_table_columns(sql))
    return schema


def validate_required_seed_values(files: list[Path], required_seed_values: dict[str, tuple[str, ...]]) -> list[str]:
    sql = "\n".join(path.read_text(encoding="utf-8") for path in files)
    errors = []
    for seed_key, values in required_seed_values.items():
        for value in values:
            if f"'{value}'" not in sql and f'"{value}"' not in sql:
                errors.append(f"missing_seed_value:{seed_key}={value}")
    return errors


def extract_created_table_columns(sql: str) -> dict[str, set[str]]:
    tables: dict[str, set[str]] = {}
    for match in re.finditer(r"create\s+table\s+if\s+not\s+exists\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\((.*?)\)\s*(?:partition\s+by\s+range\s*\([^)]+\))?\s*;", sql, flags=re.IGNORECASE | re.DOTALL):
        table_name = match.group(1)
        body = match.group(2)
        columns = set()
        for item in _split_sql_columns(body):
            token = item.strip().split(None, 1)[0].strip('"').lower() if item.strip() else ""
            if not token or token in {"primary", "unique", "constraint", "foreign", "check"}:
                continue
            columns.add(token)
        tables[table_name] = columns
    return tables


def _split_sql_columns(body: str) -> list[str]:
    items = []
    current = []
    depth = 0
    for char in body:
        if char == "(":
            depth += 1
        elif char == ")":
            depth = max(depth - 1, 0)
        if char == "," and depth == 0:
            items.append("".join(current))
            current = []
            continue
        current.append(char)
    if current:
        items.append("".join(current))
    return items


def main(env=None) -> int:
    env = os.environ if env is None else env
    result = run_migration_static_checks()
    if _get_bool(env, "MIGRATION_CHECKS_OUTPUT_JSON", False):
        print(json.dumps(result, ensure_ascii=False, sort_keys=True))
    else:
        print(f"Migration static checks passed: {result['passed']}")
        print(f"Migration files checked: {result['migration_count']}")
        if result["errors"]:
            print(f"Migration static check errors: {', '.join(result['errors'])}")
    return 0 if result["passed"] else 2


def _get_bool(env, key: str, default: bool) -> bool:
    value = env.get(key)
    if value is None or value == "":
        return default
    return str(value).strip().lower() in {"1", "true", "yes", "on"}


if __name__ == "__main__":
    raise SystemExit(main())