diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 73065b050b..96e1a5106a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -130,6 +130,60 @@ class ProxyExtrasDBManager: capture_output=True, ) + @staticmethod + def _is_permission_error(error_message: str) -> bool: + """ + Check if the error message indicates a database permission error. + + Permission errors should NOT be marked as applied, as the migration + did not actually execute successfully. + + Args: + error_message: The error message from Prisma migrate + + Returns: + bool: True if this is a permission error, False otherwise + """ + permission_patterns = [ + r"Database error code: 42501", # PostgreSQL insufficient privilege + r"must be owner of table", + r"permission denied for schema", + r"permission denied for table", + r"must be owner of schema", + ] + + for pattern in permission_patterns: + if re.search(pattern, error_message, re.IGNORECASE): + return True + return False + + @staticmethod + def _is_idempotent_error(error_message: str) -> bool: + """ + Check if the error message indicates an idempotent operation error. + + Idempotent errors (like "column already exists") mean the migration + has effectively already been applied, so it's safe to mark as applied. + + Args: + error_message: The error message from Prisma migrate + + Returns: + bool: True if this is an idempotent error, False otherwise + """ + idempotent_patterns = [ + r"already exists", + r"column .* already exists", + r"duplicate key value violates", + r"relation .* already exists", + r"constraint .* already exists", + ] + + for pattern in idempotent_patterns: + if re.search(pattern, error_message, re.IGNORECASE): + return True + return False + @staticmethod def _resolve_all_migrations( migrations_dir: str, schema_path: str, mark_all_applied: bool = True @@ -320,29 +374,79 @@ class ProxyExtrasDBManager: ) logger.info("✅ All migrations resolved.") return True - elif ( - "P3018" in e.stderr - ): # PostgreSQL error code for duplicate column - logger.info( - "Migration already exists, resolving specific migration" - ) - # Extract the migration name from the error message - migration_match = re.search( - r"Migration name: (\d+_.*)", e.stderr - ) - if migration_match: - migration_name = migration_match.group(1) - logger.info(f"Rolling back migration {migration_name}") - ProxyExtrasDBManager._roll_back_migration( - migration_name + elif "P3018" in e.stderr: + # Check if this is a permission error or idempotent error + if ProxyExtrasDBManager._is_permission_error(e.stderr): + # Permission errors should NOT be marked as applied + # Extract migration name for logging + migration_match = re.search( + r"Migration name: (\d+_.*)", e.stderr ) + migration_name = ( + migration_match.group(1) + if migration_match + else "unknown" + ) + + logger.error( + f"❌ Migration {migration_name} failed due to insufficient permissions. " + f"Please check database user privileges. Error: {e.stderr}" + ) + + # Mark as rolled back and exit with error + if migration_match: + try: + ProxyExtrasDBManager._roll_back_migration( + migration_name + ) + logger.info( + f"Migration {migration_name} marked as rolled back" + ) + except Exception as rollback_error: + logger.warning( + f"Failed to mark migration as rolled back: {rollback_error}" + ) + + # Re-raise the error to prevent silent failures + raise RuntimeError( + f"Migration failed due to permission error. Migration {migration_name} " + f"was NOT applied. Please grant necessary database permissions and retry." + ) from e + + elif ProxyExtrasDBManager._is_idempotent_error(e.stderr): + # Idempotent errors mean the migration has effectively been applied logger.info( - f"Resolving migration {migration_name} that failed due to existing columns" + "Migration failed due to idempotent error (e.g., column already exists), " + "resolving as applied" ) - ProxyExtrasDBManager._resolve_specific_migration( - migration_name + # Extract the migration name from the error message + migration_match = re.search( + r"Migration name: (\d+_.*)", e.stderr ) - logger.info("✅ Migration resolved.") + if migration_match: + migration_name = migration_match.group(1) + logger.info( + f"Rolling back migration {migration_name}" + ) + ProxyExtrasDBManager._roll_back_migration( + migration_name + ) + logger.info( + f"Resolving migration {migration_name} that failed " + f"due to existing schema objects" + ) + ProxyExtrasDBManager._resolve_specific_migration( + migration_name + ) + logger.info("✅ Migration resolved.") + else: + # Unknown P3018 error - log and re-raise for safety + logger.warning( + f"P3018 error encountered but could not classify " + f"as permission or idempotent error. " + f"Error: {e.stderr}" + ) + raise 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 index 3bcb5b25da..5714cd5c48 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,11 +1,5 @@ -import json import os import sys -import httpx -import pytest -import respx - -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../..") @@ -13,8 +7,10 @@ sys.path.insert( 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) @@ -30,3 +26,102 @@ def test_custom_prisma_dir(monkeypatch): migrations_dir = os.path.join(temp_dir, "migrations") assert os.path.exists(migrations_dir) + +class TestPermissionErrorDetection: + """Test cases for permission error detection in Prisma migrations""" + + def test_is_permission_error_postgres_42501(self): + """Test detection of PostgreSQL 42501 error code (insufficient privilege)""" + error_message = "Database error code: 42501 - permission denied for table users" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_must_be_owner(self): + """Test detection of 'must be owner of table' error""" + error_message = "ERROR: must be owner of table my_table" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_permission_denied_schema(self): + """Test detection of 'permission denied for schema' error""" + error_message = "permission denied for schema public" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_permission_denied_table(self): + """Test detection of 'permission denied for table' error""" + error_message = "permission denied for table my_table" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_must_be_owner_schema(self): + """Test detection of 'must be owner of schema' error""" + error_message = "must be owner of schema public" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_case_insensitive(self): + """Test that permission error detection is case insensitive""" + error_message = "PERMISSION DENIED FOR TABLE my_table" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + + def test_is_permission_error_negative(self): + """Test that non-permission errors are not detected as permission errors""" + error_message = "column 'id' already exists" + assert ProxyExtrasDBManager._is_permission_error(error_message) is False + + +class TestIdempotentErrorDetection: + """Test cases for idempotent error detection in Prisma migrations""" + + def test_is_idempotent_error_already_exists(self): + """Test detection of generic 'already exists' error""" + error_message = "object already exists" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + + def test_is_idempotent_error_column_already_exists(self): + """Test detection of 'column already exists' error""" + error_message = "column 'email' already exists" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + + def test_is_idempotent_error_duplicate_key(self): + """Test detection of duplicate key violation error""" + error_message = "duplicate key value violates unique constraint" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + + def test_is_idempotent_error_relation_already_exists(self): + """Test detection of 'relation already exists' error""" + error_message = "relation 'users_pkey' already exists" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + + def test_is_idempotent_error_constraint_already_exists(self): + """Test detection of 'constraint already exists' error""" + error_message = "constraint 'fk_user_id' already exists" + 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" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + + def test_is_idempotent_error_negative(self): + """Test that non-idempotent errors are not detected as idempotent errors""" + error_message = "Database error code: 42501 - permission denied" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False + + +class TestErrorClassificationPriority: + """Test cases to ensure errors are correctly classified""" + + def test_permission_error_not_classified_as_idempotent(self): + """Ensure permission errors are not mistakenly classified as idempotent""" + error_message = "Database error code: 42501 - must be owner of table users" + assert ProxyExtrasDBManager._is_permission_error(error_message) is True + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False + + def test_idempotent_error_not_classified_as_permission(self): + """Ensure idempotent errors are not mistakenly classified as permission errors""" + error_message = "column 'created_at' already exists" + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + assert ProxyExtrasDBManager._is_permission_error(error_message) is False + + def test_unknown_error_classified_as_neither(self): + """Ensure unknown errors are classified as neither permission nor idempotent""" + error_message = "connection timeout" + assert ProxyExtrasDBManager._is_permission_error(error_message) is False + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False