From 75ee0d126c957fd226bba259ee62962cc31cce2c Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 20 Jan 2026 23:23:16 +0530 Subject: [PATCH] Fix/prisma schema permission (#19391) * fix: add prisma permission issue * Add test case for prisma generate --- .../litellm_proxy_extras/utils.py | 95 +++++++++++---- litellm/proxy/db/prisma_client.py | 38 +++++- litellm/proxy/prisma_migration.py | 27 +++-- litellm/proxy/proxy_cli.py | 6 +- .../proxy/test_migration_failure_handling.py | 114 ++++++++++++++++++ 5 files changed, 238 insertions(+), 42 deletions(-) create mode 100644 tests/test_litellm/proxy/test_migration_failure_handling.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 7ffbe95be13..1aed555c5a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -18,14 +18,15 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") - def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" prisma_env = os.environ.copy() if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") + prisma_env["NPM_CONFIG_CACHE"] = os.getenv( + "NPM_CONFIG_CACHE", "/app/.cache/npm" + ) return prisma_env @@ -34,29 +35,28 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" - class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -119,7 +119,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env + env=prisma_env, ) # 3. Mark the migration as applied since it represents current state @@ -134,7 +134,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env + env=prisma_env, ) return True @@ -159,14 +159,20 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--rolled-back", + migration_name, + ], timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -178,7 +184,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -248,7 +254,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -283,7 +289,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env() + env=_get_prisma_env(), ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -313,7 +319,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -331,12 +337,18 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--applied", + migration_name, + ], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -375,7 +387,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -413,7 +425,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" @@ -509,12 +521,43 @@ class ProxyExtrasDBManager: raise else: # Use prisma db push with increased timeout - subprocess.run( - [_get_prisma_command(), "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - ) - return True + try: + subprocess.run( + [_get_prisma_command(), "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + capture_output=True, # capture output to check for errors + text=True, + env=_get_prisma_env(), + ) + return True + except subprocess.CalledProcessError as e: + if ( + "Permission denied" in e.stderr + and "schema.prisma" in e.stderr + ): + logger.warning( + f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." + ) + # Retry with --skip-generate + subprocess.run( + [ + _get_prisma_command(), + "db", + "push", + "--accept-data-loss", + "--skip-generate", + ], + timeout=60, + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + logger.info("✅ prisma db push --skip-generate completed") + return True + else: + raise e except subprocess.TimeoutExpired: logger.info(f"Attempt {attempt + 1} timed out") time.sleep(random.randrange(5, 15)) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index c9c0cfe8f68..95800e96589 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -386,11 +386,39 @@ class PrismaManager: return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) else: # Use prisma db push with increased timeout - subprocess.run( - ["prisma", "db", "push", "--accept-data-loss"], - timeout=60, - check=True, - ) + try: + subprocess.run( + ["prisma", "db", "push", "--accept-data-loss"], + timeout=60, + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as e: + if ( + "Permission denied" in e.stderr + and "schema.prisma" in e.stderr + ): + verbose_proxy_logger.warning( + f"Permission denied during prisma generate: {e.stderr}. Retrying with --skip-generate..." + ) + # Retry with --skip-generate + subprocess.run( + [ + "prisma", + "db", + "push", + "--accept-data-loss", + "--skip-generate", + ], + timeout=60, + check=True, + capture_output=True, + text=True, + ) + return True + else: + raise e return True except subprocess.TimeoutExpired: verbose_proxy_logger.warning(f"Attempt {attempt + 1} timed out") diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 251d1e56287..2fe12b1439c 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -15,14 +15,21 @@ from litellm.proxy.proxy_cli import run_server # Call the Click command with standalone_mode=False run_server(["--skip_server_startup"], standalone_mode=False) -# run prisma generate +# Run prisma generate verbose_proxy_logger.info("Running 'prisma generate'...") -result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) -verbose_proxy_logger.info(f"'prisma generate' stdout: {result.stdout}") # Log stdout -exit_code = result.returncode - -if exit_code != 0: - verbose_proxy_logger.info(f"'prisma generate' failed with exit code {exit_code}.") - verbose_proxy_logger.error( - f"'prisma generate' stderr: {result.stderr}" - ) # Log stderr +try: + result = subprocess.run(["prisma", "generate"], capture_output=True, text=True) + if result.returncode != 0: + if "Permission denied" in result.stderr: + verbose_proxy_logger.warning( + f"Permission denied during 'prisma generate'. Skipping generation, assuming client is pre-generated. Error: {result.stderr}" + ) + else: + verbose_proxy_logger.info( + f"'prisma generate' failed with exit code {result.returncode}." + ) + verbose_proxy_logger.error( + f"'prisma generate' stderr: {result.stderr}" + ) # Log stderr +except Exception as e: + verbose_proxy_logger.error(f"Error running prisma generate: {e}") diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 2059246674b..028aad9a2e1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -797,7 +797,11 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - PrismaManager.setup_database(use_migrate=not use_prisma_db_push) + if not PrismaManager.setup_database( + use_migrate=not use_prisma_db_push + ): + print("LiteLLM: Database setup failed. Exiting...") # noqa + sys.exit(1) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa diff --git a/tests/test_litellm/proxy/test_migration_failure_handling.py b/tests/test_litellm/proxy/test_migration_failure_handling.py new file mode 100644 index 00000000000..426049007b1 --- /dev/null +++ b/tests/test_litellm/proxy/test_migration_failure_handling.py @@ -0,0 +1,114 @@ +import sys +import os +import subprocess +from unittest.mock import MagicMock, patch +from click.testing import CliRunner + +# Add parent directory to path to allow importing litellm +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.prisma_client import PrismaManager +from litellm.proxy.proxy_cli import run_server + + +class TestMigrationFailureHandling: + @patch("subprocess.run") + def test_prisma_client_permission_error_retry(self, mock_subprocess_run): + """ + Regression Test: Verifies that PrismaManager.setup_database + catches PermissionError during 'prisma db push' and retries with '--skip-generate'. + """ + # Mock behavior: + # call 1: raises CalledProcessError with "Permission denied" and "schema.prisma" + # call 2 (retry): succeeds + + error_output = "Error: Permission denied writing to ... schema.prisma" + + mock_process_error = subprocess.CalledProcessError( + returncode=1, cmd=["prisma", "db", "push"], stderr=error_output + ) + + mock_subprocess_run.side_effect = [ + mock_process_error, # 1st attempt fails with permission error + MagicMock(returncode=0), # 2nd attempt (retry) succeeds + ] + + # Ensure we run the 'db push' path (use_migrate=False) + # We also need to mock should_update_prisma_schema to return True + + with patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", + return_value=True, + ): + # Run setup_database with use_migrate=False to trigger 'prisma db push' path + result = PrismaManager.setup_database(use_migrate=False) + + # Assert success + assert result is True + + # Verify calls + assert mock_subprocess_run.call_count == 2 + + # Check 1st call arguments (standard push) + args1, _ = mock_subprocess_run.call_args_list[0] + assert "push" in args1[0] + assert "--skip-generate" not in args1[0] + + # Check 2nd call arguments (retry with skip-generate) + args2, _ = mock_subprocess_run.call_args_list[1] + assert "push" in args2[0] + assert "--skip-generate" in args2[0] + + def test_proxy_cli_exit_on_migration_fail(self): + """ + Regression Test: Verifies that proxy_cli.run_server exits with NON-ZERO status + if PrismaManager.setup_database returns False. + """ + runner = CliRunner() + + # Mock setup_database to return False (Simulating failure) + # Mock should_update_prisma_schema to return True (Ensure we hit the DB setup logic) + with patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database", + return_value=False, + ), patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", + return_value=True, + ): + # Mock dependencies to prevent actual server startup and handle imports + mock_app = MagicMock() + mock_proxy_config = MagicMock() + + # Patch sys.modules to prevent ImportErrors for proxy_server + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, ProxyConfig=mock_proxy_config + ) + }, + ): + with patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "app", + "host": "localhost", + "port": 8000, + } + + # Set DATABASE_URL to trigger DB logic + with patch.dict( + os.environ, + {"DATABASE_URL": "postgresql://user:pass@localhost:5432/db"}, + ): + # Execute: Run server with --local and --skip_server_startup + result = runner.invoke( + run_server, ["--local", "--skip_server_startup"] + ) + + # Assert: Exit code should be non-zero (failure) + assert ( + result.exit_code != 0 + ), f"Expected non-zero exit code, got {result.exit_code}. Output: {result.output}" + assert "Database setup failed. Exiting..." in result.output