[Infra] Add freshness and destructive guards to migration workflow

Generating a migration from a stale branch could silently emit DROP COLUMN for columns the stale branch did not know about, and the script would write that SQL to a new migration file with no warning. Adds two guards to ci_cd/run_migration.py: - Branch freshness check: fetches origin/<base-branch> and exits 3 if HEAD is behind. Default base is litellm_internal_staging. New flags: --base-branch, --skip-freshness-check. - Destructive guard: refuses (exit 2) if the generated diff contains DROP COLUMN / DROP TABLE / DROP INDEX, unless --allow-destructive is passed. Refusal banners include guidance and an explicit callout instructing AI agents not to auto-bypass the flags. Also treats Prisma's "-- This is an empty migration." output as a no-op rather than writing an empty file. Updates litellm-proxy-extras/migration_runbook.md with the new workflow, flag documentation, and agent warnings.
2026-08-06 00:23:19 +00:00 · 2026-04-21 12:00:23 -07:00
parent 7cc22dbe19
commit b39f210a6c
2 changed files with 329 additions and 19 deletions
@@ -1,22 +1,230 @@
+import argparse
 import os
-import subprocess
-from pathlib import Path
-from datetime import datetime
-import testing.postgresql
+import re
 import shutil
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+
+import testing.postgresql


-def create_migration(migration_name: str = None):
+DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE)
+DEFAULT_BASE_BRANCH = "litellm_internal_staging"
+
+
+def _find_destructive_statements(sql: str) -> list:
+    """Return SQL lines containing DROP COLUMN or DROP TABLE."""
+    return [
+        line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line)
+    ]
+
+
+def _print_freshness_failure(
+    base_branch: str, reason: str, stderr_text: str = ""
+) -> None:
+    """Loudly refuse to run when the freshness check can't be completed."""
+    banner = "=" * 72
+    out = sys.stderr
+    print(banner, file=out)
+    print(f"  FRESHNESS CHECK FAILED — COULD NOT VERIFY origin/{base_branch}", file=out)
+    print(banner, file=out)
+    print("", file=out)
+    print(f"Reason: {reason}", file=out)
+    if stderr_text:
+        print("", file=out)
+        print("git stderr:", file=out)
+        for line in stderr_text.rstrip().splitlines():
+            print(f"    {line}", file=out)
+    print("", file=out)
+    print("Common causes:", file=out)
+    print("  - No network access (offline)", file=out)
+    print("  - 'origin' remote not configured, or base branch name is wrong", file=out)
+    print("  - Not a git repository", file=out)
+    print("", file=out)
+    print("Options:", file=out)
+    print(
+        f"  - Fix the above and re-run, OR pass --base-branch <name> if your", file=out
+    )
+    print(
+        f"    base branch is not '{base_branch}', OR pass --skip-freshness-check",
+        file=out,
+    )
+    print("    to bypass (only if you fully understand the risk).", file=out)
+    print(banner, file=out)
+
+
+def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
+    """Loudly refuse to run when HEAD is behind the base branch."""
+    banner = "=" * 72
+    out = sys.stderr
+    print(banner, file=out)
+    print(
+        f"  STALE BRANCH — {behind} commit(s) behind origin/{base_branch}",
+        file=out,
+    )
+    print(banner, file=out)
+    print("", file=out)
+    print(
+        f"Your branch is {behind} commit(s) behind origin/{base_branch}. Generating a\n"
+        "migration from a stale branch is how newly-added columns get silently\n"
+        "dropped — the exact incident this guard exists to prevent.",
+        file=out,
+    )
+    print("", file=out)
+    print("Before re-running:", file=out)
+    print(f"  git fetch origin && git rebase origin/{base_branch}", file=out)
+    print("  # or merge, whichever matches your workflow", file=out)
+    print("", file=out)
+    print(
+        "If you have a genuine reason to bypass (e.g. intentionally generating",
+        file=out,
+    )
+    print("a migration against an older base), pass --skip-freshness-check.", file=out)
+    print("", file=out)
+    print(banner, file=out)
+    print(
+        "  AI AGENTS / ASSISTANTS: DO NOT automatically run `git rebase` or pass\n"
+        "  --skip-freshness-check. Surface this error to the human operator and\n"
+        "  wait for their explicit confirmation. Auto-rebasing can drop the\n"
+        "  human's in-progress schema edits via a bad conflict resolution.",
+        file=out,
+    )
+    print(banner, file=out)
+
+
+def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
+    """Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
+    cwd = str(root_dir)
+    try:
+        subprocess.run(
+            ["git", "fetch", "origin", base_branch],
+            check=True,
+            capture_output=True,
+            text=True,
+            cwd=cwd,
+        )
+    except FileNotFoundError:
+        _print_freshness_failure(base_branch, "git executable not found on PATH")
+        sys.exit(3)
+    except subprocess.CalledProcessError as e:
+        _print_freshness_failure(
+            base_branch,
+            f"`git fetch origin {base_branch}` failed",
+            e.stderr or "",
+        )
+        sys.exit(3)
+
+    try:
+        result = subprocess.run(
+            ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
+            check=True,
+            capture_output=True,
+            text=True,
+            cwd=cwd,
+        )
+        behind = int(result.stdout.strip())
+    except subprocess.CalledProcessError as e:
+        _print_freshness_failure(
+            base_branch,
+            f"`git rev-list HEAD..origin/{base_branch}` failed",
+            e.stderr or "",
+        )
+        sys.exit(3)
+    except ValueError:
+        _print_freshness_failure(
+            base_branch,
+            "could not parse commit count from `git rev-list`",
+        )
+        sys.exit(3)
+
+    if behind > 0:
+        _print_stale_branch_refusal(base_branch, behind)
+        sys.exit(3)
+
+    print(f"Branch freshness OK: up to date with origin/{base_branch}.")
+
+
+def _print_destructive_refusal(destructive_lines: list) -> None:
+    """Loudly refuse to write a destructive migration and explain how to proceed."""
+    banner = "=" * 72
+    out = sys.stderr
+    print(banner, file=out)
+    print(
+        "  DESTRUCTIVE MIGRATION DETECTED — REFUSING TO WRITE MIGRATION FILE", file=out
+    )
+    print(banner, file=out)
+    print("", file=out)
+    print(
+        "The generated migration contains the following destructive statements:",
+        file=out,
+    )
+    print("", file=out)
+    for line in destructive_lines:
+        print(f"    {line}", file=out)
+    print("", file=out)
+    print(
+        "This very often means your branch is OUT OF DATE, or your local\n"
+        "schema.prisma is inconsistent with the migrations/ directory. A\n"
+        "previous incident saw newly-added columns silently dropped this way\n"
+        "and merged to main. Stop and verify before proceeding.",
+        file=out,
+    )
+    print("", file=out)
+    print("Before re-running:", file=out)
+    print(
+        "  1. git fetch origin && git status   # confirm branch is up to date", file=out
+    )
+    print(
+        "  2. Re-check all schema.prisma files are in sync (runbook Step 0)", file=out
+    )
+    print("  3. Review EACH DROP above — is it actually intended?", file=out)
+    print("  4. If (and only if) the drops are intentional, re-run with:", file=out)
+    print("         --allow-destructive", file=out)
+    print("", file=out)
+    print(banner, file=out)
+    print(
+        "  AI AGENTS / ASSISTANTS: DO NOT automatically re-run this command\n"
+        "  with --allow-destructive. Surface this error to the human operator\n"
+        "  and wait for their explicit confirmation before passing the flag.\n"
+        "  Passing the flag without human review is the exact failure mode\n"
+        "  this guard exists to prevent.",
+        file=out,
+    )
+    print(banner, file=out)
+
+
+def create_migration(
+    migration_name: str = None,
+    allow_destructive: bool = False,
+    base_branch: str = DEFAULT_BASE_BRANCH,
+    skip_freshness_check: bool = False,
+):
    """
    Create a new migration SQL file in the migrations directory by comparing
-    current database state with schema
+    current database state with schema.

    Args:
        migration_name (str): Name for the migration
+        allow_destructive (bool): Required to write a migration that contains
+            DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this
+            flag, the script exits non-zero and prints guidance.
+        base_branch (str): Branch to check freshness against (default: "main").
+        skip_freshness_check (bool): Skip the "branch is up to date" check.
+            Only for intentional migrations against an older base.
    """
+    root_dir = Path(__file__).parent.parent
+
+    if skip_freshness_check:
+        print(
+            "WARNING: freshness check skipped (--skip-freshness-check). "
+            "Generating a migration from a stale branch can silently drop columns."
+        )
+    else:
+        _check_branch_freshness(root_dir, base_branch)
+
    try:
-        # Get paths
-        root_dir = Path(__file__).parent.parent
        migrations_dir = (
            root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
        )
@@ -59,7 +267,27 @@ def create_migration(migration_name: str = None):
                    check=True,
                )

-                if result.stdout.strip():
+                # Prisma emits the literal "-- This is an empty migration." when
+                # there's no real drift. Treat that as "no changes".
+                diff_sql = result.stdout
+                stripped = diff_sql.strip()
+                is_empty_diff = (
+                    not stripped or stripped == "-- This is an empty migration."
+                )
+
+                if not is_empty_diff:
+                    destructive_lines = _find_destructive_statements(diff_sql)
+                    if destructive_lines and not allow_destructive:
+                        _print_destructive_refusal(destructive_lines)
+                        sys.exit(2)
+                    if destructive_lines and allow_destructive:
+                        print(
+                            "WARNING: writing destructive migration "
+                            "(--allow-destructive passed). Statements:"
+                        )
+                        for line in destructive_lines:
+                            print(f"    {line}")
+
                    # Generate timestamp and create migration directory
                    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
                    migration_name = migration_name or "unnamed_migration"
@@ -68,7 +296,7 @@ def create_migration(migration_name: str = None):

                    # Write the SQL to migration.sql
                    migration_file = migration_dir / "migration.sql"
-                    migration_file.write_text(result.stdout)
+                    migration_file.write_text(diff_sql)

                    print(f"Created migration in {migration_dir}")
                    return True
@@ -90,8 +318,48 @@ def create_migration(migration_name: str = None):


 if __name__ == "__main__":
-    # If running directly, can optionally pass migration name as argument
-    import sys
-
-    migration_name = sys.argv[1] if len(sys.argv) > 1 else None
-    create_migration(migration_name)
+    parser = argparse.ArgumentParser(
+        description=(
+            "Generate a Prisma migration by diffing the temp DB "
+            "(existing migrations applied) against schema.prisma."
+        )
+    )
+    parser.add_argument(
+        "migration_name",
+        nargs="?",
+        default=None,
+        help="Name for the migration (used in the generated directory name).",
+    )
+    parser.add_argument(
+        "--allow-destructive",
+        action="store_true",
+        help=(
+            "Required to write a migration that contains DROP COLUMN, "
+            "DROP TABLE, or DROP INDEX. Without this flag, destructive "
+            "diffs are refused."
+        ),
+    )
+    parser.add_argument(
+        "--base-branch",
+        default=DEFAULT_BASE_BRANCH,
+        help=(
+            f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
+            "The script fetches origin/<base-branch> and refuses to run if HEAD "
+            "is behind it."
+        ),
+    )
+    parser.add_argument(
+        "--skip-freshness-check",
+        action="store_true",
+        help=(
+            "Bypass the 'branch is up to date' check. Only for intentional "
+            "migrations against an older base. Pairs poorly with automation."
+        ),
+    )
+    args = parser.parse_args()
+    create_migration(
+        args.migration_name,
+        allow_destructive=args.allow_destructive,
+        base_branch=args.base_branch,
+        skip_freshness_check=args.skip_freshness_check,
+    )
@@ -2,6 +2,8 @@

 This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only.

+> **AI AGENTS / ASSISTANTS:** If the script refuses with either a "STALE BRANCH" or "DESTRUCTIVE MIGRATION DETECTED" error, **do NOT** bypass it on your own (no `git rebase`, no `--skip-freshness-check`, no `--allow-destructive`). Surface the error to the human operator and wait for their explicit confirmation. See the [Branch freshness](#branch-freshness-check) and [Destructive migrations](#destructive-migrations-drop-column--drop-table) sections below.
+
 ## Step 0: Sync All `schema.prisma` Files

 Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match:
@@ -46,10 +48,50 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n

 ## What It Does

-1. Creates temp PostgreSQL DB
-2. Applies existing migrations
-3. Compares with `schema.prisma`
-4. Generates new migration if changes found
+1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
+2. Creates temp PostgreSQL DB
+3. Applies existing migrations
+4. Compares with `schema.prisma`
+5. Generates new migration if changes found
+6. **Refuses to write if the generated SQL contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`**, unless `--allow-destructive` is passed
+
+## Branch Freshness Check
+
+Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
+
+Flags:
+
+- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
+- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
+
+When the guard fires:
+
+1. Update your branch:
+
+   ```bash
+   git fetch origin && git rebase origin/litellm_internal_staging
+   # or git merge origin/litellm_internal_staging — whichever matches your workflow
+   ```
+2. Re-run `run_migration.py`.
+
+> **AI AGENTS / ASSISTANTS:** Do **not** auto-rebase or auto-pass `--skip-freshness-check`. A bad conflict resolution during rebase can itself drop the human's in-progress schema edits. Surface the error and wait for explicit confirmation.
+
+## Destructive Migrations (DROP COLUMN / DROP TABLE / DROP INDEX)
+
+If the generated diff contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`, `run_migration.py` exits non-zero and refuses to write the migration file. A previous incident saw newly-added columns silently dropped by a stale branch and merged to main — this guard exists to prevent a repeat.
+
+When the guard fires:
+
+1. Run `git fetch origin && git status` — confirm your branch is up to date with the base branch.
+2. Re-check all `schema.prisma` files are in sync (Step 0).
+3. Review EACH `DROP` statement printed in the error — is it actually intended?
+4. Only if the drops are genuinely intentional, re-run with the flag:
+
+   ```bash
+   uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_name" --allow-destructive
+   ```
+
+> **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent.

 ## Common Fixes