From 120e3d6d3dd141cd13f16422ca132b35b7297a0d Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 21 Jul 2026 06:36:59 +0000 Subject: [PATCH] fix(proxy): handle migration and credential DB failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/utils.py | 19 ++++- .../tests/test_setup_database_fail_fast.py | 76 +++++++++++++++++++ .../proxy/credential_endpoints/endpoints.py | 9 ++- .../credential_endpoints/test_endpoints.py | 23 ++++++ 4 files changed, 123 insertions(+), 4 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f7c0bed0800..87511a8599e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -530,6 +530,7 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) timeout_attempts = 0 + baseline_attempts = 0 baseline_created = False recovered_migrations: frozenset[str] = frozenset() try: @@ -566,8 +567,22 @@ class ProxyExtrasDBManager: "migration ledger after baseline creation." ) from e logger.info("Schema exists but no migrations ledger — creating baseline") - ProxyExtrasDBManager._create_baseline_migration(schema_path) - baseline_created = True + baseline_attempts += 1 + try: + baseline_created = ProxyExtrasDBManager._create_baseline_migration(schema_path) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as baseline_error: + raise RuntimeError( + "Failed to create the database migration baseline. " + "Check database connectivity and permissions." + ) from baseline_error + if not baseline_created and baseline_attempts >= 4: + raise RuntimeError( + "Database migration baseline creation failed after 4 attempts. " + "Check database connectivity and load." + ) continue if "P3009" in stderr: diff --git a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py index fb83872bf6c..d968ea63efa 100644 --- a/litellm-proxy-extras/tests/test_setup_database_fail_fast.py +++ b/litellm-proxy-extras/tests/test_setup_database_fail_fast.py @@ -255,3 +255,79 @@ def test_v2_recovers_multiple_idempotent_migrations_in_one_invocation(monkeypatc assert run.call_count == 6 assert roll_back.call_args_list == [call(name) for name in migration_names] assert resolve.call_args_list == [call(name) for name in migration_names] + + +def test_v2_retries_baseline_after_timeout(monkeypatch, tmp_path): + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + p3005 = subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate deploy", + stderr="Error: P3005\ndatabase schema is not empty", + output="", + ) + + class FakeResult: + stdout = "Applied migration.\n" + stderr = "" + + run = MagicMock(side_effect=(p3005, p3005, FakeResult())) + create_baseline = MagicMock(side_effect=(False, True)) + monkeypatch.setattr("subprocess.run", run) + monkeypatch.setattr(ProxyExtrasDBManager, "_create_baseline_migration", create_baseline) + + ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert ok is True + assert run.call_count == 3 + assert create_baseline.call_count == 2 + + +def test_v2_baseline_timeouts_are_bounded(monkeypatch, tmp_path): + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + p3005 = subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate deploy", + stderr="Error: P3005\ndatabase schema is not empty", + output="", + ) + run = MagicMock(side_effect=(p3005, p3005, p3005, p3005)) + create_baseline = MagicMock(return_value=False) + monkeypatch.setattr("subprocess.run", run) + monkeypatch.setattr(ProxyExtrasDBManager, "_create_baseline_migration", create_baseline) + + with pytest.raises(RuntimeError, match="baseline creation failed after 4 attempts"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + + assert run.call_count == 4 + assert create_baseline.call_count == 4 + + +def test_v2_baseline_subprocess_failure_raises_runtime_error(monkeypatch, tmp_path): + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + (tmp_path / "schema.prisma").write_text("// stub") + p3005 = subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate deploy", + stderr="Error: P3005\ndatabase schema is not empty", + output="", + ) + baseline_error = subprocess.CalledProcessError( + returncode=1, + cmd="prisma migrate resolve --applied 0_init", + stderr="permission denied", + output="", + ) + monkeypatch.setattr("subprocess.run", MagicMock(side_effect=p3005)) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_create_baseline_migration", + MagicMock(side_effect=baseline_error), + ) + + with pytest.raises(RuntimeError, match="Failed to create the database migration baseline"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 0db1d6bddb4..1bd772d1957 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -6,7 +6,9 @@ import asyncio from dataclasses import dataclass from typing import TYPE_CHECKING, Optional +import httpx from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response +from prisma.errors import PrismaError if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -319,7 +321,7 @@ async def get_credentials( return {"success": True, "credentials": masked_credentials} except HTTPException: raise - except Exception as e: + except (PrismaError, httpx.HTTPError) as e: return handle_exception_on_proxy(e) @@ -595,7 +597,10 @@ async def update_credential( ) ) credentials_repository = CredentialsRepository(prisma_client) - existing = await credentials_repository.find_by_name(credential_name) + try: + existing = await credentials_repository.find_by_name(credential_name) + except Exception as e: + return handle_exception_on_proxy(e) await _authorize_credential_patch( credential_name=credential_name, patch=credential, diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index 5337ef7e1f8..4b3f6e6f002 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -14,6 +14,7 @@ import sys import pytest from fastapi import HTTPException +from prisma.errors import ClientNotConnectedError from unittest.mock import AsyncMock, MagicMock sys.path.insert(0, os.path.abspath("../../../..")) @@ -431,6 +432,28 @@ async def test_team_admin_cannot_replay_stale_global_scope(_connected_db, _patch _connected_db.update_by_name.assert_not_awaited() +@pytest.mark.asyncio +async def test_update_credential_handles_database_lookup_failure(_connected_db): + db_error = ClientNotConnectedError() + _connected_db.find_by_name = AsyncMock(side_effect=db_error) + _connected_db.update_by_name = AsyncMock() + + result = await endpoints.update_credential( + request=MagicMock(), + fastapi_response=MagicMock(), + credential=UpdateCredentialItem( + credential_info={"access": {"global": True}}, + ), + credential_name="dest", + user_api_key_dict=_admin(), + ) + + assert result.code == "500" + assert result.message == str(db_error) + _connected_db.find_by_name.assert_awaited_once_with("dest") + _connected_db.update_by_name.assert_not_awaited() + + @pytest.mark.asyncio async def test_provider_credential_patch_forbidden_for_non_admin(_connected_db, monkeypatch): """A team-admin (or any non-admin) cannot PATCH a non-logging credential.