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/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:

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)

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,

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.