scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

Extract _deploy_with_idempotent_resolution loop for P3009/P3018

Browse source 
The previous approach resolved one idempotent migration per outer retry
attempt, exhausting all 5 retries when a DB had 4+ idempotent failures
(e.g., tables/columns already created by the old force-apply logic).

Now _deploy_with_idempotent_resolution loops internally: deploy → detect
idempotent P3009/P3018 → rollback + resolve → re-deploy, until all
idempotent migrations are resolved or a non-recoverable error is hit.
This keeps outer retries available for transient errors (timeouts, etc).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-12 11:29:42 -07:00
parent 347090d288
commit d41e5e29ce
2 changed files with 202 additions and 200 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

295
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -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(

107
tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
View file

@ -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"]