diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 58c4e174867..3446aa36600 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -283,6 +283,137 @@ class ProxyExtrasDBManager: 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 # 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: """ @@ -306,82 +437,11 @@ class ProxyExtrasDBManager: 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") + 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 not migration_match: - # Cannot identify which migration failed — fail fast - raise RuntimeError( - f"Migration failed (P3009) but could not extract migration name. " - f"Manual intervention required. Error: {e.stderr}" - ) from e - - 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. Re-deploying remaining migrations..." - ) - # Re-run deploy to apply any migrations after the resolved one - 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 - else: - 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, fix the issue, " - f"and restart.\n" - f"Original error: {e.stderr}" - ) from e - elif ( + if ( "P3005" in e.stderr and "database schema is not empty" in e.stderr ): @@ -395,96 +455,15 @@ class ProxyExtrasDBManager: ProxyExtrasDBManager._mark_all_migrations_applied( migrations_dir ) - # Now run prisma migrate deploy to apply any truly pending migrations + # Now run deploy with resolution for any 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}" - ) + ProxyExtrasDBManager._deploy_with_idempotent_resolution() logger.info("✅ All migrations 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( 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 af136e8f96f..a8eff29b17b 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -214,46 +214,65 @@ class TestMarkAllMigrationsApplied: ProxyExtrasDBManager._mark_all_migrations_applied("/fake/migrations/dir") -class TestSetupDatabaseFailFast: - """Test that setup_database fails fast on non-recoverable migration errors""" +class TestDeployWithIdempotentResolution: + """Test _deploy_with_idempotent_resolution loops through multiple idempotent failures""" - @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""" + 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="", ) - # First call (migrate deploy) raises P3009; subsequent calls (roll_back) succeed mock_run.side_effect = [deploy_error, MagicMock(returncode=0)] with pytest.raises(RuntimeError, match="requires manual intervention"): - ProxyExtrasDBManager.setup_database(use_migrate=True) + ProxyExtrasDBManager._deploy_with_idempotent_resolution() - # Verify rollback was called (second subprocess call) + # 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.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_unmatched_regex_raises_runtime_error( - self, mock_run, mock_dir, mock_getcwd, mock_chdir - ): - """P3009 with unparseable migration name should fail fast, not silently retry""" + def test_p3009_unmatched_regex_raises_runtime_error(self, mock_run): + """P3009 with unparseable migration name should fail fast""" error = subprocess.CalledProcessError( 1, "prisma", @@ -263,28 +282,19 @@ class TestSetupDatabaseFailFast: mock_run.side_effect = error with pytest.raises(RuntimeError, match="could not extract migration name"): - ProxyExtrasDBManager.setup_database(use_migrate=True) + ProxyExtrasDBManager._deploy_with_idempotent_resolution() - # Should fail on first attempt, not retry assert mock_run.call_count == 1 - @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_idempotent_redeploys_remaining_migrations( - self, mock_run, mock_dir, mock_getcwd, mock_chdir - ): - """P3009 with idempotent error should resolve the failed migration then re-deploy""" + 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="", ) - # Call sequence: deploy (fails P3009), roll_back, resolve, re-deploy (succeeds) mock_run.side_effect = [ deploy_error, MagicMock(returncode=0), # roll_back @@ -292,30 +302,43 @@ class TestSetupDatabaseFailFast: MagicMock(stdout="All migrations applied", returncode=0), # re-deploy ] - result = ProxyExtrasDBManager.setup_database(use_migrate=True) + ProxyExtrasDBManager._deploy_with_idempotent_resolution() - assert result is True assert mock_run.call_count == 4 - # Last call should be prisma migrate deploy (re-deploy) 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_does_not_call_resolve_all( - self, mock_run, mock_dir, mock_getcwd, mock_chdir - ): + 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 - # 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"]