From 93b6df96e0ffc5397c75ea44fc1430aad584f458 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 26 Apr 2025 12:05:06 -0700 Subject: [PATCH] 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 https://github.com/BerriAI/litellm/issues/9885 * fix(utils.py): give helpful warning message when permission denied error raised in fs --- .gitignore | 1 + docs/my-website/docs/proxy/config_settings.md | 1 + .../litellm_proxy_extras/utils.py | 139 ++++++++++++++++-- litellm/proxy/db/prisma_client.py | 6 +- .../test_litellm_proxy_extras_utils.py | 32 ++++ 5 files changed, 158 insertions(+), 21 deletions(-) create mode 100644 tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py diff --git a/.gitignore b/.gitignore index e8c18bed4cb..2ba3272259a 100644 --- a/.gitignore +++ b/.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/* diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 1e3c800b037..0f88c79cf2d 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index e771e48e453..21c9131887b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -2,8 +2,10 @@ import glob 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 @@ class ProxyExtrasDBManager: # 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 @@ class ProxyExtrasDBManager: ) 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 @@ class ProxyExtrasDBManager: ) @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 @@ class ProxyExtrasDBManager: ) @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 @@ class ProxyExtrasDBManager: 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 @@ class ProxyExtrasDBManager: 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 ( diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 339117cb272..c4edc702b32 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -137,7 +137,6 @@ class PrismaManager: 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 @@ class PrismaManager: 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( diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py new file mode 100644 index 00000000000..3bcb5b25da8 --- /dev/null +++ b/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) +