diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl new file mode 100644 index 00000000000..9a5c185de28 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz new file mode 100644 index 00000000000..3e4be95b519 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55-py3-none-any.whl new file mode 100644 index 00000000000..f44d0744ef2 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55.tar.gz new file mode 100644 index 00000000000..e0e471d0f27 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.55.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql deleted file mode 100644 index 7b3e6d089ec..00000000000 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309115809_add_missing_indexes/migration.sql +++ /dev/null @@ -1,13 +0,0 @@ --- SkipTransactionBlock - --- Drop invalid indexes left behind by failed CONCURRENTLY builds -DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_VerificationToken_key_alias_idx"; - --- CreateIndex -CREATE INDEX CONCURRENTLY "LiteLLM_VerificationToken_key_alias_idx" ON "LiteLLM_VerificationToken"("key_alias"); - --- Drop invalid indexes left behind by failed CONCURRENTLY builds -DROP INDEX CONCURRENTLY IF EXISTS "LiteLLM_SpendLogs_user_startTime_idx"; - --- CreateIndex -CREATE INDEX CONCURRENTLY "LiteLLM_SpendLogs_user_startTime_idx" ON "LiteLLM_SpendLogs"("user", "startTime"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql new file mode 100644 index 00000000000..5ab834695b8 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql @@ -0,0 +1,11 @@ +-- DropIndex +DROP INDEX "LiteLLM_MCPServerTable_approval_status_idx"; + +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "approval_status", +DROP COLUMN "review_notes", +DROP COLUMN "reviewed_at", +DROP COLUMN "source_url", +DROP COLUMN "submitted_at", +DROP COLUMN "submitted_by"; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d5d17b2bcec..8d4bdffb2dd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -388,9 +388,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -556,9 +553,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f3155722187..30bbdc9a57b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -5,7 +5,6 @@ import shutil import subprocess import time -from datetime import datetime from pathlib import Path from typing import Optional @@ -231,7 +230,6 @@ def _is_idempotent_error(error_message: str) -> bool: idempotent_patterns = [ r"already exists", r"column .* already exists", - r"duplicate key value violates", r"relation .* already exists", r"constraint .* already exists", r"does not exist", @@ -244,97 +242,18 @@ def _is_idempotent_error(error_message: str) -> bool: return False @staticmethod - def _resolve_all_migrations( - migrations_dir: str, schema_path: str, mark_all_applied: bool = True - ): + def _mark_all_migrations_applied(migrations_dir: 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") - if not database_url: - logger.error("DATABASE_URL not set") - return - - 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( - [ - _get_prisma_command(), - "migrate", - "diff", - "--from-url", - database_url, - "--to-schema-datamodel", - schema_path, - "--script", - ], - check=True, - timeout=60, - stdout=f, - env=_get_prisma_env(), - ) - 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.") + Mark all existing migrations as applied in the _prisma_migrations table. - # 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}") + Used after creating a baseline migration for an existing database (P3005), + so that Prisma knows these migrations have already been reflected in the schema. - # 2. Run prisma db execute to apply the migration - try: - logger.info("Running prisma db execute to apply the migration diff...") - result = subprocess.run( - [ - _get_prisma_command(), - "db", - "execute", - "--file", - str(diff_sql_path), - "--schema", - schema_path, - ], - timeout=60, - check=True, - capture_output=True, - text=True, - env=_get_prisma_env(), - ) - 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 - if not mark_all_applied: - return + This does NOT generate or apply any schema diffs — it only updates migration + tracking state. + """ migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir) - logger.info(f"Resolving {len(migration_names)} migrations") + logger.info(f"Marking {len(migration_names)} migrations as applied") for migration_name in migration_names: try: logger.info(f"Resolving migration: {migration_name}") @@ -354,10 +273,145 @@ def _resolve_all_migrations( ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: - if "is already recorded as applied in the database." not in e.stderr: - logger.warning( - f"Failed to resolve migration {migration_name}: {e.stderr}" + if "is already recorded as applied in the database." in e.stderr: + logger.debug( + f"Migration {migration_name} already recorded as applied" ) + else: + raise RuntimeError( + f"Failed to mark migration {migration_name} as applied: {e.stderr}" + ) from e + + @staticmethod + def _resolve_failed_migration(e: subprocess.CalledProcessError): + """ + Handle a failed migration (P3009 or P3018) by resolving idempotent errors + or raising for non-recoverable errors. + + Raises: + RuntimeError: If the error is non-idempotent, a permission error, or + the migration name cannot be extracted. + """ + stderr = e.stderr + + # Determine error code and extract migration name + if "P3009" in stderr: + migration_match = re.search(r"`(\d+_.*)` migration", stderr) + if not migration_match: + raise RuntimeError( + f"Migration failed (P3009) but could not extract migration name. " + f"Manual intervention required. Error: {stderr}" + ) from e + migration_name = migration_match.group(1) + elif "P3018" in stderr: + if ProxyExtrasDBManager._is_permission_error(stderr): + migration_match = re.search( + r"Migration name: (\d+_.*)", 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: {stderr}" + ) + 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}" + ) + 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 + + migration_match = re.search(r"Migration name: (\d+_.*)", stderr) + if not migration_match: + raise RuntimeError( + f"Migration failed (P3018) but could not extract migration name. " + f"Manual intervention required. Error: {stderr}" + ) from e + migration_name = migration_match.group(1) + else: + raise e # Not a P3009/P3018 — let outer handler deal with it + + # Check if idempotent — if not, fail fast + if not ProxyExtrasDBManager._is_idempotent_error(stderr): + logger.error( + f"❌ Migration {migration_name} failed with a non-idempotent error. " + f"This requires manual intervention. Error: {stderr}" + ) + 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}" + ) + raise RuntimeError( + f"Migration {migration_name} failed and requires manual intervention. " + f"Please inspect the migration and database state, fix the issue, " + f"and restart.\n" + f"Original error: {stderr}" + ) from e + + # Idempotent error — resolve and continue + logger.info( + f"Migration {migration_name} failed due to idempotent error " + f"(e.g., column already exists), resolving as applied" + ) + ProxyExtrasDBManager._roll_back_migration(migration_name) + ProxyExtrasDBManager._resolve_specific_migration(migration_name) + logger.info(f"✅ Migration {migration_name} resolved.") + + @staticmethod + def _deploy_with_idempotent_resolution(max_resolutions: int = 150): + """ + Run prisma migrate deploy, automatically resolving idempotent failures + (P3009/P3018 with 'already exists' etc.) in a loop. + + Stops when deploy succeeds or a non-idempotent error is encountered. + The max_resolutions cap prevents infinite loops if something goes wrong. + + Raises: + RuntimeError: On non-recoverable migration errors. + subprocess.CalledProcessError: On unexpected Prisma errors. + """ + for i in range(max_resolutions): + try: + result = subprocess.run( + [_get_prisma_command(), "migrate", "deploy"], + timeout=60, + check=True, + capture_output=True, + text=True, + env=_get_prisma_env(), + ) + logger.info(f"prisma migrate deploy stdout: {result.stdout}") + logger.info("✅ prisma migrate deploy completed") + return + except subprocess.CalledProcessError as e: + logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}") + if "P3009" in e.stderr or "P3018" in e.stderr: + # Raises RuntimeError for non-recoverable errors, + # returns normally for resolved idempotent errors + ProxyExtrasDBManager._resolve_failed_migration(e) + logger.info("Re-deploying remaining migrations...") + continue + else: + raise + + raise RuntimeError( + f"Exceeded maximum idempotent resolutions ({max_resolutions}). " + f"This likely indicates a deeper issue with migration state." + ) @staticmethod def setup_database(use_migrate: bool = False) -> bool: @@ -373,7 +427,7 @@ def setup_database(use_migrate: bool = False) -> bool: bool: True if setup was successful, False otherwise """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" - for attempt in range(4): + for attempt in range(5): original_dir = os.getcwd() migrations_dir = ProxyExtrasDBManager._get_prisma_dir() os.chdir(migrations_dir) @@ -382,72 +436,11 @@ def setup_database(use_migrate: bool = False) -> bool: if use_migrate: logger.info("Running prisma migrate deploy") try: - # Set migrations directory for Prisma - result = subprocess.run( - [_get_prisma_command(), "migrate", "deploy"], - timeout=60, - check=True, - capture_output=True, - text=True, - env=_get_prisma_env(), - ) - logger.info(f"prisma migrate deploy stdout: {result.stdout}") - - logger.info("prisma migrate deploy completed") - - # Run sanity check to ensure DB matches schema - logger.info("Running post-migration sanity check...") - ProxyExtrasDBManager._resolve_all_migrations( - migrations_dir, schema_path, mark_all_applied=False - ) - logger.info("✅ Post-migration sanity check completed") + ProxyExtrasDBManager._deploy_with_idempotent_resolution() return True except subprocess.CalledProcessError as e: logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}") - if "P3009" in e.stderr: - # Extract the failed migration name from the error message - migration_match = re.search( - r"`(\d+_.*)` migration", e.stderr - ) - if migration_match: - failed_migration = migration_match.group(1) - 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 ( + if ( "P3005" in e.stderr and "database schema is not empty" in e.stderr ): @@ -456,86 +449,17 @@ def setup_database(use_migrate: bool = False) -> bool: ) ProxyExtrasDBManager._create_baseline_migration(schema_path) logger.info( - "Baseline migration created, resolving all migrations" + "Baseline migration created, marking all existing migrations as applied. " + "This is the standard Prisma baselining approach for databases " + "previously managed by db push." ) - ProxyExtrasDBManager._resolve_all_migrations( - migrations_dir, schema_path + ProxyExtrasDBManager._mark_all_migrations_applied( + migrations_dir ) - logger.info("✅ All migrations resolved.") + logger.info("✅ All migrations marked as applied.") return True - 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( - "Migration failed due to idempotent error (e.g., column already exists), " - "resolving as applied" - ) - # 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 - ) - 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: + raise else: # Use prisma db push with increased timeout subprocess.run( @@ -548,7 +472,7 @@ def setup_database(use_migrate: bool = False) -> bool: logger.info(f"Attempt {attempt + 1} timed out") time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: - attempts_left = 3 - attempt + attempts_left = 4 - attempt retry_msg = ( f" Retrying... ({attempts_left} attempts left)" if attempts_left > 0 diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index ef80f092f1b..6f1cf9b9d0b 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.53" +version = "0.4.55" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.53" +version = "0.4.55" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3af72d65b56..721c3e404d2 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -397,9 +397,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -565,9 +562,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request diff --git a/pyproject.toml b/pyproject.toml index dd8747b6649..caa11e44b53 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "^0.4.53", optional = true} +litellm-proxy-extras = {version = "^0.4.55", optional = true} rich = {version = "^13.7.1", optional = true} litellm-enterprise = {version = "^0.1.33", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index ccbfa281d91..9095d05ea08 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.53 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.55 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index d5d17b2bcec..8d4bdffb2dd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -388,9 +388,6 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) - - // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (...) ORDER BY "public"."LiteLLM_VerificationToken"."key_alias" ASC - @@index([key_alias]) } model LiteLLM_JWTKeyMapping { @@ -556,9 +553,6 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) - - // SELECT ... FROM "LiteLLM_SpendLogs" WHERE ("startTime" >= $1 AND "startTime" <= $2 AND "user" = $3) GROUP BY ... - @@index([user, startTime]) } // View spend, model, api_key per request 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 597c0845d43..38890ff8c5e 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,5 +1,9 @@ import os +import subprocess import sys +from unittest.mock import MagicMock, patch + +import pytest sys.path.insert( 0, @@ -82,10 +86,10 @@ def test_is_idempotent_error_column_already_exists(self): 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""" + def test_duplicate_key_violation_not_idempotent(self): + """Duplicate key violations (e.g., CREATE INDEX with duplicate data) are NOT idempotent""" error_message = "duplicate key value violates unique constraint" - assert ProxyExtrasDBManager._is_idempotent_error(error_message) is True + assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False def test_is_idempotent_error_relation_already_exists(self): """Test detection of 'relation already exists' error""" @@ -138,3 +142,203 @@ def test_unknown_error_classified_as_neither(self): error_message = "connection timeout" assert ProxyExtrasDBManager._is_permission_error(error_message) is False assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False + + +class TestMarkAllMigrationsApplied: + """Test that _mark_all_migrations_applied only marks migrations without applying diffs""" + + @patch("litellm_proxy_extras.utils.subprocess.run") + @patch.object( + ProxyExtrasDBManager, + "_get_migration_names", + return_value=["20250326162113_baseline", "20250329084805_new_cron_job_table"], + ) + def test_marks_each_migration_as_applied(self, mock_get_names, mock_run): + """Verify each migration is marked as applied via prisma migrate resolve""" + ProxyExtrasDBManager._mark_all_migrations_applied("/fake/migrations/dir") + + assert mock_run.call_count == 2 + for call_args in mock_run.call_args_list: + cmd = call_args[0][0] + assert "migrate" in cmd + assert "resolve" in cmd + assert "--applied" in cmd + + @patch("litellm_proxy_extras.utils.subprocess.run") + @patch.object( + ProxyExtrasDBManager, + "_get_migration_names", + return_value=["20250326162113_baseline"], + ) + def test_does_not_generate_or_apply_diffs(self, mock_get_names, mock_run): + """Verify no diff generation or db execute commands are run""" + ProxyExtrasDBManager._mark_all_migrations_applied("/fake/migrations/dir") + + for call_args in mock_run.call_args_list: + cmd = call_args[0][0] + # Should never run diff or db execute + assert "diff" not in cmd + assert "execute" not in cmd + assert "push" not in cmd + + @patch("litellm_proxy_extras.utils.subprocess.run") + @patch.object( + ProxyExtrasDBManager, + "_get_migration_names", + return_value=["20250326162113_baseline"], + ) + def test_skips_already_applied_migration(self, mock_get_names, mock_run): + """Verify already-applied migrations are silently skipped""" + mock_run.side_effect = subprocess.CalledProcessError( + 1, + "prisma", + stderr="Migration `20250326162113_baseline` is already recorded as applied in the database.", + ) + # Should not raise + ProxyExtrasDBManager._mark_all_migrations_applied("/fake/migrations/dir") + + @patch("litellm_proxy_extras.utils.subprocess.run") + @patch.object( + ProxyExtrasDBManager, + "_get_migration_names", + return_value=["20250326162113_baseline"], + ) + def test_raises_on_unexpected_error(self, mock_get_names, mock_run): + """Verify non-'already applied' errors are propagated, not silently swallowed""" + mock_run.side_effect = subprocess.CalledProcessError( + 1, + "prisma", + stderr="connection refused", + ) + with pytest.raises(RuntimeError, match="Failed to mark migration"): + ProxyExtrasDBManager._mark_all_migrations_applied("/fake/migrations/dir") + + +class TestDeployWithIdempotentResolution: + """Test _deploy_with_idempotent_resolution loops through multiple idempotent failures""" + + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_resolves_multiple_idempotent_migrations_in_one_pass(self, mock_run): + """Multiple P3018 idempotent errors should all be resolved without consuming outer retries""" + p3018_error_1 = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3018\nMigration name: 20251113000000_add_project_table\nERROR: relation \"LiteLLM_ProjectTable\" already exists", + output="", + ) + p3018_error_2 = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3018\nMigration name: 20251113000001_add_project_fields\nERROR: column \"description\" already exists", + output="", + ) + # deploy fails, rollback, resolve, deploy fails again, rollback, resolve, deploy succeeds + mock_run.side_effect = [ + p3018_error_1, + MagicMock(returncode=0), # roll_back + MagicMock(returncode=0), # resolve + p3018_error_2, + MagicMock(returncode=0), # roll_back + MagicMock(returncode=0), # resolve + MagicMock(stdout="All migrations applied", returncode=0), # final deploy + ] + + ProxyExtrasDBManager._deploy_with_idempotent_resolution() + + assert mock_run.call_count == 7 + # First, fourth, and seventh calls should be deploy + for idx in [0, 3, 6]: + cmd = mock_run.call_args_list[idx][0][0] + assert cmd == ["prisma", "migrate", "deploy"] + + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_p3009_non_idempotent_raises_runtime_error(self, mock_run): + """P3009 with non-idempotent error should raise RuntimeError""" + deploy_error = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3009: migrate found failed migrations in the target database, `20250329084805_new_cron_job_table` migration. Error: syntax error at or near 'ALTR'", + output="", + ) + mock_run.side_effect = [deploy_error, MagicMock(returncode=0)] + + with pytest.raises(RuntimeError, match="requires manual intervention"): + ProxyExtrasDBManager._deploy_with_idempotent_resolution() + + # deploy + rollback + assert mock_run.call_count == 2 + rollback_cmd = mock_run.call_args_list[1][0][0] + assert "--rolled-back" in rollback_cmd + + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_p3009_unmatched_regex_raises_runtime_error(self, mock_run): + """P3009 with unparseable migration name should fail fast""" + error = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3009: migrate found failed migrations in the target database, unexpected format", + output="", + ) + mock_run.side_effect = error + + with pytest.raises(RuntimeError, match="could not extract migration name"): + ProxyExtrasDBManager._deploy_with_idempotent_resolution() + + assert mock_run.call_count == 1 + + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_p3009_idempotent_redeploys(self, mock_run): + """P3009 with idempotent error should resolve then re-deploy""" + deploy_error = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3009: migrate found failed migrations in the target database, `20250329084805_new_cron_job_table` migration. Error: column 'status' already exists", + output="", + ) + mock_run.side_effect = [ + deploy_error, + MagicMock(returncode=0), # roll_back + MagicMock(returncode=0), # resolve + MagicMock(stdout="All migrations applied", returncode=0), # re-deploy + ] + + ProxyExtrasDBManager._deploy_with_idempotent_resolution() + + assert mock_run.call_count == 4 + last_cmd = mock_run.call_args_list[3][0][0] + assert last_cmd == ["prisma", "migrate", "deploy"] + + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_p3018_permission_error_raises(self, mock_run): + """P3018 with permission error should raise RuntimeError""" + deploy_error = subprocess.CalledProcessError( + 1, + "prisma", + stderr="P3018\nMigration name: 20251113000000_add_project_table\nDatabase error code: 42501\npermission denied for table users", + output="", + ) + mock_run.side_effect = [deploy_error, MagicMock(returncode=0)] + + with pytest.raises(RuntimeError, match="permission error"): + ProxyExtrasDBManager._deploy_with_idempotent_resolution() + + +class TestSetupDatabase: + """Test setup_database integration with deploy and resolution""" + + @patch("litellm_proxy_extras.utils.os.chdir") + @patch("litellm_proxy_extras.utils.os.getcwd", return_value="/original") + @patch.object( + ProxyExtrasDBManager, "_get_prisma_dir", return_value="/fake/prisma/dir" + ) + @patch("litellm_proxy_extras.utils.subprocess.run") + def test_successful_deploy(self, mock_run, mock_dir, mock_getcwd, mock_chdir): + """After successful prisma migrate deploy, no diff/resolve should be called""" + mock_run.return_value = MagicMock(stdout="All migrations applied", returncode=0) + + result = ProxyExtrasDBManager.setup_database(use_migrate=True) + + assert result is True + assert mock_run.call_count == 1 + cmd = mock_run.call_args[0][0] + assert cmd == ["prisma", "migrate", "deploy"]