Merge c56e00afcb into 44d84360fb

ci_cd/run_migration.py

@@ -4,6 +4,7 @@ import re
 import shutil
 import subprocess
 import sys
+import tempfile
 from datetime import datetime
 from pathlib import Path
 
@@ -195,6 +196,25 @@ def _print_destructive_refusal(destructive_lines: list) -> None:
     print(banner, file=out)
 
 
+def _prepare_isolated_scratch(
+    schema_path: Path, migrations_dir: Path
+) -> tuple[Path, Path]:
+    """Copy schema.prisma and existing migrations into a fresh temp directory.
+
+    Prisma resolves the migrations directory as a "migrations" folder next to
+    the schema file, so both must be copied together. A unique ``mkdtemp`` root
+    is used (never ``schema_path.parent``) so the scratch copy can't collide
+    with the tracked repo-root ``migrations/`` folder and delete it on cleanup.
+
+    Returns ``(temp_root, temp_schema_path)``; the caller removes ``temp_root``.
+    """
+    temp_root = Path(tempfile.mkdtemp(prefix="litellm_migrations_"))
+    temp_schema_path = temp_root / schema_path.name
+    shutil.copy2(schema_path, temp_schema_path)
+    shutil.copytree(migrations_dir, temp_root / "migrations")
+    return temp_root, temp_schema_path
+
+
 def create_migration(
     migration_name: str = None,
     allow_destructive: bool = False,
@@ -235,19 +255,15 @@ def create_migration(
         with testing.postgresql.Postgresql() as postgresql:
             db_url = postgresql.url()
 
-            # Create temporary migrations directory next to schema.prisma
-            temp_migrations_dir = schema_path.parent / "migrations"
+            temp_root, temp_schema_path = _prepare_isolated_scratch(
+                schema_path, migrations_dir
+            )
 
             try:
-                # Copy existing migrations to temp directory
-                if temp_migrations_dir.exists():
-                    shutil.rmtree(temp_migrations_dir)
-                shutil.copytree(migrations_dir, temp_migrations_dir)
-
                 # Apply existing migrations to temp database
                 os.environ["DATABASE_URL"] = db_url
                 subprocess.run(
-                    ["prisma", "migrate", "deploy", "--schema", str(schema_path)],
+                    ["prisma", "migrate", "deploy", "--schema", str(temp_schema_path)],
                     check=True,
                 )
 
@@ -260,7 +276,7 @@ def create_migration(
                         "--from-url",
                         db_url,
                         "--to-schema-datamodel",
-                        str(schema_path),
+                        str(temp_schema_path),
                         "--script",
                     ],
                     capture_output=True,
@@ -306,9 +322,8 @@ def create_migration(
                     return False
 
             finally:
-                # Clean up: remove temporary migrations directory
-                if temp_migrations_dir.exists():
-                    shutil.rmtree(temp_migrations_dir)
+                # Clean up: remove temporary directory
+                shutil.rmtree(temp_root, ignore_errors=True)
 
     except subprocess.CalledProcessError as e:
         print(f"Error generating migration: {e.stderr}")

tests/proxy_migration_tests/test_run_migration.py

@@ -0,0 +1,70 @@
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+if "testing.postgresql" not in sys.modules:
+    _testing = types.ModuleType("testing")
+    _testing_postgresql = types.ModuleType("testing.postgresql")
+    _testing.postgresql = _testing_postgresql
+    sys.modules.setdefault("testing", _testing)
+    sys.modules.setdefault("testing.postgresql", _testing_postgresql)
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+_SPEC = importlib.util.spec_from_file_location(
+    "run_migration", REPO_ROOT / "ci_cd" / "run_migration.py"
+)
+run_migration = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(run_migration)
+
+
+def test_prepare_isolated_scratch_does_not_touch_repo_root_migrations(tmp_path):
+    """Regression for #33135: scratch dir must not collide with repo migrations/.
+
+    ``create_migration`` used ``schema_path.parent / "migrations"`` as its
+    scratch directory and ``rmtree``'d it, wiping the tracked repo-root
+    ``migrations/`` folder (Dockerfile, run.py). The scratch copy must live in
+    an isolated temp dir instead.
+    """
+    repo_root = tmp_path / "repo"
+    repo_root.mkdir()
+
+    schema_path = repo_root / "schema.prisma"
+    schema_path.write_text("// schema")
+
+    # tracked repo-root migrations/ that must survive untouched
+    tracked_migrations = repo_root / "migrations"
+    tracked_migrations.mkdir()
+    dockerfile = tracked_migrations / "Dockerfile"
+    dockerfile.write_text("FROM scratch")
+    run_py = tracked_migrations / "run.py"
+    run_py.write_text("print('hi')")
+
+    # the real migrations source that gets copied into the scratch dir
+    source_migrations = repo_root / "litellm_proxy_extras" / "migrations"
+    source_migrations.mkdir(parents=True)
+    (source_migrations / "20240101000000_init").mkdir()
+    (source_migrations / "20240101000000_init" / "migration.sql").write_text(
+        "SELECT 1;"
+    )
+
+    temp_root, temp_schema_path = run_migration._prepare_isolated_scratch(
+        schema_path, source_migrations
+    )
+    try:
+        assert temp_root.resolve() != tracked_migrations.resolve()
+        assert repo_root.resolve() not in temp_root.resolve().parents
+
+        assert temp_schema_path.read_text() == "// schema"
+        assert (temp_root / "migrations" / "20240101000000_init" / "migration.sql").read_text() == "SELECT 1;"
+
+        assert dockerfile.read_text() == "FROM scratch"
+        assert run_py.read_text() == "print('hi')"
+        assert sorted(p.name for p in tracked_migrations.iterdir()) == [
+            "Dockerfile",
+            "run.py",
+        ]
+    finally:
+        import shutil
+
+        shutil.rmtree(temp_root, ignore_errors=True)