Prisma Migrate - support setting custom migration dir (#10336)

* build(litellm-proxy-extras/utils.py): correctly generate baseline migration for non-empty db

* fix(litellm-proxy-extras/utils.py): Fix issue in migration, where if a migration fails during baselining, all are still marked as applied

* fix(prisma_client.py): don't pass separate schema.prisma to litellm-proxy-extras

use the one in litellm-proxy-extras

* fix(litellm-proxy-extras/utils.py): support passing custom dir for baselining db in read-only fs

Fixes #9885

* fix(utils.py): give helpful warning message when permission denied error raised in fs

Author: krrishdholakia

.gitignore

@@ -88,3 +88,4 @@ litellm/proxy/migrations/*config.yaml
 litellm/proxy/migrations/*
 config.yaml
 tests/litellm/litellm_core_utils/llm_cost_calc/log.txt
+tests/test_custom_dir/*

docs/my-website/docs/proxy/config_settings.md

@@ -442,6 +442,7 @@ router_settings:
 | LITELLM_EMAIL | Email associated with LiteLLM account
 | LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES | Maximum retries for parallel requests in LiteLLM
 | LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRY_TIMEOUT | Timeout for retries of parallel requests in LiteLLM
+| LITELLM_MIGRATION_DIR | Custom migrations directory for prisma migrations, used for baselining db in read-only file systems.
 | LITELLM_HOSTED_UI | URL of the hosted UI for LiteLLM
 | LITELLM_LICENSE | License key for LiteLLM usage
 | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -2,8 +2,10 @@
 import os
 import random
 import re
+import shutil
 import subprocess
 import time
+from datetime import datetime
 from pathlib import Path
 from typing import Optional
 
@@ -19,9 +21,30 @@ def str_to_bool(value: Optional[str]) -> bool:
 class ProxyExtrasDBManager:
     @staticmethod
     def _get_prisma_dir() -> str:
-        """Get the path to the migrations directory"""
-        migrations_dir = os.path.dirname(__file__)
-        return migrations_dir
+        """
+        Get the path to the migrations directory
+
+        Set os.environ["LITELLM_MIGRATION_DIR"] to a custom migrations directory, to support baselining db in read-only fs.
+        """
+        custom_migrations_dir = os.getenv("LITELLM_MIGRATION_DIR")
+        pkg_migrations_dir = os.path.dirname(__file__)
+        if custom_migrations_dir:
+            # If migrations_dir exists, copy contents
+            if os.path.exists(custom_migrations_dir):
+                # Copy contents instead of directory itself
+                for item in os.listdir(pkg_migrations_dir):
+                    src_path = os.path.join(pkg_migrations_dir, item)
+                    dst_path = os.path.join(custom_migrations_dir, item)
+                    if os.path.isdir(src_path):
+                        shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
+                    else:
+                        shutil.copy2(src_path, dst_path)
+            else:
+                # If directory doesn't exist, create it and copy everything
+                shutil.copytree(pkg_migrations_dir, custom_migrations_dir)
+            return custom_migrations_dir
+
+        return pkg_migrations_dir
 
     @staticmethod
     def _create_baseline_migration(schema_path: str) -> bool:
@@ -33,27 +56,29 @@ def _create_baseline_migration(schema_path: str) -> bool:
         # Create migrations/0_init directory
         init_dir.mkdir(parents=True, exist_ok=True)
 
-        # Generate migration SQL file
-        migration_file = init_dir / "migration.sql"
+        database_url = os.getenv("DATABASE_URL")
 
         try:
-            # Generate migration diff with increased timeout
+            # 1. Generate migration SQL file by comparing empty state to current db state
+            logger.info("Generating baseline migration...")
+            migration_file = init_dir / "migration.sql"
             subprocess.run(
                 [
                     "prisma",
                     "migrate",
                     "diff",
                     "--from-empty",
-                    "--to-schema-datamodel",
-                    str(schema_path),
+                    "--to-url",
+                    database_url,
                     "--script",
                 ],
                 stdout=open(migration_file, "w"),
                 check=True,
                 timeout=30,
-            )  # 30 second timeout
+            )
 
-            # Mark migration as applied with increased timeout
+            # 3. Mark the migration as applied since it represents current state
+            logger.info("Marking baseline migration as applied...")
             subprocess.run(
                 [
                     "prisma",
@@ -73,8 +98,10 @@ def _create_baseline_migration(schema_path: str) -> bool:
             )
             return False
         except subprocess.CalledProcessError as e:
-            logger.warning(f"Error creating baseline migration: {e}")
-            return False
+            logger.warning(
+                f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}"
+            )
+            raise e
 
     @staticmethod
     def _get_migration_names(migrations_dir: str) -> list:
@@ -104,8 +131,85 @@ def _resolve_specific_migration(migration_name: str):
         )
 
     @staticmethod
-    def _resolve_all_migrations(migrations_dir: str):
-        """Mark all existing migrations as applied"""
+    def _resolve_all_migrations(migrations_dir: str, schema_path: str):
+        """
+        1. Compare the current database state to schema.prisma and generate a migration for the diff.
+        2. Run prisma migrate deploy to apply any pending migrations.
+        3. Mark all existing migrations as applied.
+        """
+        database_url = os.getenv("DATABASE_URL")
+        diff_dir = (
+            Path(migrations_dir)
+            / "migrations"
+            / f"{datetime.now().strftime('%Y%m%d%H%M%S')}_baseline_diff"
+        )
+        try:
+            diff_dir.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            if "Permission denied" in str(e):
+                logger.warning(
+                    f"Permission denied - {e}\nunable to baseline db. Set LITELLM_MIGRATION_DIR environment variable to a writable directory to enable migrations."
+                )
+                return
+            raise e
+        diff_sql_path = diff_dir / "migration.sql"
+
+        # 1. Generate migration SQL for the diff between DB and schema
+        try:
+            logger.info("Generating migration diff between DB and schema.prisma...")
+            with open(diff_sql_path, "w") as f:
+                subprocess.run(
+                    [
+                        "prisma",
+                        "migrate",
+                        "diff",
+                        "--from-url",
+                        database_url,
+                        "--to-schema-datamodel",
+                        schema_path,
+                        "--script",
+                    ],
+                    check=True,
+                    timeout=60,
+                    stdout=f,
+                )
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to generate migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff generation timed out.")
+
+        # check if the migration was created
+        if not diff_sql_path.exists():
+            logger.warning("Migration diff was not created")
+            return
+        logger.info(f"Migration diff created at {diff_sql_path}")
+
+        # 2. Run prisma db execute to apply the migration
+        try:
+            logger.info("Running prisma db execute to apply the migration diff...")
+            result = subprocess.run(
+                [
+                    "prisma",
+                    "db",
+                    "execute",
+                    "--file",
+                    str(diff_sql_path),
+                    "--schema",
+                    schema_path,
+                ],
+                timeout=60,
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.info(f"prisma db execute stdout: {result.stdout}")
+            logger.info("✅ Migration diff applied successfully")
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to apply migration diff: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            logger.warning("Migration diff application timed out.")
+
+        # 3. Mark all migrations as applied
         migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
         logger.info(f"Resolving {len(migration_names)} migrations")
         for migration_name in migration_names:
@@ -126,7 +230,7 @@ def _resolve_all_migrations(migrations_dir: str):
                     )
 
     @staticmethod
-    def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
+    def setup_database(use_migrate: bool = False) -> bool:
         """
         Set up the database using either prisma migrate or prisma db push
         Uses migrations from litellm-proxy-extras package
@@ -138,6 +242,7 @@ def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
         Returns:
             bool: True if setup was successful, False otherwise
         """
+        schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
         use_migrate = str_to_bool(os.getenv("USE_PRISMA_MIGRATE")) or use_migrate
         for attempt in range(4):
             original_dir = os.getcwd()
@@ -200,7 +305,9 @@ def setup_database(schema_path: str, use_migrate: bool = False) -> bool:
                             logger.info(
                                 "Baseline migration created, resolving all migrations"
                             )
-                            ProxyExtrasDBManager._resolve_all_migrations(migrations_dir)
+                            ProxyExtrasDBManager._resolve_all_migrations(
+                                migrations_dir, schema_path
+                            )
                             logger.info("✅ All migrations resolved.")
                             return True
                         elif (

litellm/proxy/db/prisma_client.py

@@ -137,7 +137,6 @@ def setup_database(use_migrate: bool = False) -> bool:
         for attempt in range(4):
             original_dir = os.getcwd()
             prisma_dir = PrismaManager._get_prisma_dir()
-            schema_path = prisma_dir + "/schema.prisma"
             os.chdir(prisma_dir)
             try:
                 if use_migrate:
@@ -150,11 +149,8 @@ def setup_database(use_migrate: bool = False) -> bool:
                         return False
 
                     prisma_dir = PrismaManager._get_prisma_dir()
-                    schema_path = prisma_dir + "/schema.prisma"
 
-                    return ProxyExtrasDBManager.setup_database(
-                        schema_path=schema_path, use_migrate=use_migrate
-                    )
+                    return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate)
                 else:
                     # Use prisma db push with increased timeout
                     subprocess.run(

tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

@@ -0,0 +1,32 @@
+import json
+import os
+import sys
+import httpx
+import pytest
+import respx
+
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+
+from litellm_proxy_extras.utils import ProxyExtrasDBManager
+
+def test_custom_prisma_dir(monkeypatch):
+    import tempfile
+    # create a temp directory
+    temp_dir = tempfile.mkdtemp()
+    monkeypatch.setenv("LITELLM_MIGRATION_DIR", temp_dir)
+
+    ## Check if the prisma dir is the temp directory
+    assert ProxyExtrasDBManager._get_prisma_dir() == temp_dir
+
+    ## Check if the schema.prisma file is in the temp directory
+    schema_path = os.path.join(temp_dir, "schema.prisma")
+    assert os.path.exists(schema_path)
+
+    ## Check if the migrations dir is in the temp directory
+    migrations_dir = os.path.join(temp_dir, "migrations")
+    assert os.path.exists(migrations_dir)
+