diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 7ffbe95be13..f3155722187 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 @@ def _create_baseline_migration(schema_path: str) -> bool: 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 @@ def _create_baseline_migration(schema_path: str) -> bool: ], check=True, timeout=30, - env=prisma_env + env=prisma_env, ) return True @@ -159,14 +159,20 @@ def _get_migration_names(migrations_dir: str) -> list: @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 @@ def _resolve_specific_migration(migration_name: str): timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -228,6 +234,8 @@ def _is_idempotent_error(error_message: str) -> bool: r"duplicate key value violates", r"relation .* already exists", r"constraint .* already exists", + r"does not exist", + r"Can't drop database.* because it doesn't exist", ] for pattern in idempotent_patterns: @@ -248,7 +256,7 @@ def _resolve_all_migrations( if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -283,7 +291,7 @@ def _resolve_all_migrations( 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 +321,7 @@ def _resolve_all_migrations( 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 +339,18 @@ def _resolve_all_migrations( 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 +389,7 @@ def setup_database(use_migrate: bool = False) -> bool: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -397,27 +411,42 @@ def setup_database(use_migrate: bool = False) -> bool: ) if migration_match: failed_migration = migration_match.group(1) - logger.info( - f"Found failed migration: {failed_migration}, marking as rolled back" - ) - # Mark the failed migration as rolled back - subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--rolled-back", - failed_migration, - ], - timeout=60, - check=True, - capture_output=True, - text=True, - env=_get_prisma_env() - ) - logger.info( - f"✅ Migration {failed_migration} marked as rolled back... retrying" - ) + if ProxyExtrasDBManager._is_idempotent_error(e.stderr): + logger.info( + f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied" + ) + ProxyExtrasDBManager._roll_back_migration( + failed_migration + ) + ProxyExtrasDBManager._resolve_specific_migration( + failed_migration + ) + logger.info( + f"✅ Migration {failed_migration} resolved." + ) + return True + else: + logger.info( + f"Found failed migration: {failed_migration}, marking as rolled back" + ) + # Mark the failed migration as rolled back + subprocess.run( + [ + _get_prisma_command(), + "migrate", + "resolve", + "--rolled-back", + failed_migration, + ], + timeout=60, + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + logger.info( + f"✅ Migration {failed_migration} marked as rolled back... retrying" + ) elif ( "P3005" in e.stderr and "database schema is not empty" in e.stderr diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 5714cd5c487..6787d0cf72d 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -94,6 +94,11 @@ def test_is_idempotent_error_constraint_already_exists(self): error_message = "constraint 'fk_user_id' already exists" assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + def test_is_idempotent_error_does_not_exist(self): + """Test detection of 'does not exist' error""" + error_message = "ERROR: index 'idx' does not exist" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + def test_is_idempotent_error_case_insensitive(self): """Test that idempotent error detection is case insensitive""" error_message = "COLUMN 'ID' ALREADY EXISTS"