diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f3155722187..8d527351297 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -5,7 +5,6 @@ import re import shutil import subprocess import time -from datetime import datetime from pathlib import Path from typing import Optional @@ -244,97 +243,18 @@ class ProxyExtrasDBManager: 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. + Mark all existing migrations as applied in the _prisma_migrations table. + + Used after creating a baseline migration for an existing database (P3005), + so that Prisma knows these migrations have already been reflected in the schema. + + This does NOT generate or apply any schema diffs — it only updates migration + tracking state. """ - 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.") - - # 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( - [ - _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 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}") @@ -392,15 +312,7 @@ class ProxyExtrasDBManager: 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") + logger.info("✅ prisma migrate deploy completed") return True except subprocess.CalledProcessError as e: logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}") @@ -426,27 +338,29 @@ class ProxyExtrasDBManager: ) 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" + logger.error( + f"❌ Migration {failed_migration} failed with a non-idempotent error. " + f"This requires manual intervention. Error: {e.stderr}" ) + # Mark as rolled back so the migration can be retried after manual fix + try: + ProxyExtrasDBManager._roll_back_migration( + failed_migration + ) + logger.info( + f"Migration {failed_migration} 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_migration} failed and requires manual intervention. " + f"Please inspect the migration and database state, then either:\n" + f" - Fix the issue and restart, or\n" + f" - Run: prisma migrate resolve --applied {failed_migration}\n" + f"Original error: {e.stderr}" + ) from e elif ( "P3005" in e.stderr and "database schema is not empty" in e.stderr @@ -456,12 +370,27 @@ class ProxyExtrasDBManager: ) ProxyExtrasDBManager._create_baseline_migration(schema_path) logger.info( - "Baseline migration created, resolving all migrations" + "Baseline migration created, marking all existing migrations as applied" ) - ProxyExtrasDBManager._resolve_all_migrations( - migrations_dir, schema_path + ProxyExtrasDBManager._mark_all_migrations_applied( + migrations_dir ) - logger.info("✅ All migrations resolved.") + # Now run prisma migrate deploy to apply any truly pending migrations + logger.info( + "Running prisma migrate deploy for any pending migrations..." + ) + 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("✅ All migrations applied.") return True elif "P3018" in e.stderr: # Check if this is a permission error or idempotent error 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..cf084cc54e1 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, @@ -138,3 +142,102 @@ class TestErrorClassificationPriority: 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") + + +class TestSetupDatabaseFailFast: + """Test that setup_database fails fast on non-recoverable migration errors""" + + @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_p3009_non_idempotent_raises_runtime_error( + self, mock_run, mock_dir, mock_getcwd, mock_chdir + ): + """P3009 with non-idempotent error should raise RuntimeError, not silently retry""" + 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 = error + + with pytest.raises(RuntimeError, match="requires manual intervention"): + ProxyExtrasDBManager.setup_database(use_migrate=True) + + @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_does_not_call_resolve_all( + 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 + # Only one subprocess call: prisma migrate deploy + assert mock_run.call_count == 1 + cmd = mock_run.call_args[0][0] + assert cmd == ["prisma", "migrate", "deploy"]