Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_passthrough_live_credentials

litellm/proxy/proxy_server.py

@@ -8664,48 +8664,56 @@ class ProxyStartupEvent:
         - Sets up prisma client
         - Adds necessary views to proxy
         """
+        connected_client: PrismaClient | None = None
         try:
-            prisma_client: PrismaClient | None = None
-            if database_url is not None:
-                try:
-                    prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj)
-                except Exception as e:
-                    raise e
+            if database_url is None:
+                return None
 
-                try:
-                    await prisma_client.connect()
-                except Exception as e:
-                    if "P3018" in str(e) or "P3009" in str(e):
-                        verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED")
-                        verbose_proxy_logger.debug("Your database is in a 'dirty' state.")
-                        verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied <migration_name>'")
-                    raise e
+            prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj)
 
-                ## Start RDS IAM token refresh background task if enabled ##
-                # This proactively refreshes IAM tokens before they expire,
-                # preventing the 15-minute connection failure bug (#16220)
-                if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"):
-                    await prisma_client.db.start_token_refresh_task()
+            try:
+                await prisma_client.connect()
+            except Exception as e:
+                if "P3018" in str(e) or "P3009" in str(e):
+                    verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED")
+                    verbose_proxy_logger.debug("Your database is in a 'dirty' state.")
+                    verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied <migration_name>'")
+                raise e
 
-                ## Add necessary views to proxy ##
-                asyncio.create_task(
-                    prisma_client.check_view_exists()
-                )  # check if all necessary views exist. Don't block execution
+            connected_client = prisma_client
 
-                asyncio.create_task(
-                    prisma_client._set_spend_logs_row_count_in_proxy_state()
-                )  # set the spend logs row count in proxy state. Don't block execution
+            ## Start RDS IAM token refresh background task if enabled ##
+            # This proactively refreshes IAM tokens before they expire,
+            # preventing the 15-minute connection failure bug (#16220)
+            if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"):
+                await prisma_client.db.start_token_refresh_task()
 
-                # run a health check to ensure the DB is ready
-                if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True:
-                    await prisma_client.health_check()
+            ## Add necessary views to proxy ##
+            asyncio.create_task(
+                prisma_client.check_view_exists()
+            )  # check if all necessary views exist. Don't block execution
+
+            asyncio.create_task(
+                prisma_client._set_spend_logs_row_count_in_proxy_state()
+            )  # set the spend logs row count in proxy state. Don't block execution
+
+            if hasattr(prisma_client, "start_db_health_watchdog_task"):
+                await prisma_client.start_db_health_watchdog_task()
+
+            # run a health check to ensure the DB is ready
+            if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True:
+                await prisma_client.health_check()
 
-                if hasattr(prisma_client, "start_db_health_watchdog_task"):
-                    await prisma_client.start_db_health_watchdog_task()
             return prisma_client
         except Exception as e:
             PrismaDBExceptionHandler.handle_db_exception(e)
-            return None
+            if connected_client is not None:
+                verbose_proxy_logger.warning(
+                    "Retaining the connected Prisma client after a post-connect startup step failed: %s. "
+                    "The DB health watchdog keeps probing and reconnects once the database recovers.",
+                    e,
+                )
+            return connected_client
 
     @classmethod
     def _init_dd_tracer(cls):

litellm/proxy/utils.py

@@ -4269,7 +4269,7 @@ class PrismaClient:
             import traceback
 
             error_msg: Final = f"LiteLLM Prisma Client Exception connect(): {e}"
-            print_verbose(error_msg)
+            verbose_proxy_logger.warning(error_msg)
             error_traceback: Final = error_msg + "\n" + traceback.format_exc()
             end_time: Final = time.time()
             _duration: Final = end_time - start_time
@@ -4987,8 +4987,8 @@ class PrismaClient:
         except Exception as e:
             import traceback
 
-            error_msg: Final = f"LiteLLM Prisma Client Exception disconnect(): {e}"
-            print_verbose(error_msg)
+            error_msg: Final = f"LiteLLM Prisma Client Exception health_check(): {e}"
+            verbose_proxy_logger.warning(error_msg)
             error_traceback: Final = error_msg + "\n" + traceback.format_exc()
             end_time: Final = time.time()
             _duration: Final = end_time - start_time

tests/test_litellm/proxy/test_proxy_server.py

@@ -11018,3 +11018,123 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog):
         ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={})
 
     assert MOCK_TESTING_CONFIG_KEY not in caplog.text
+
+
+def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
+    client = MagicMock()
+    client.connect = AsyncMock(side_effect=connect_error)
+    client.db.start_token_refresh_task = AsyncMock()
+    client.check_view_exists = AsyncMock()
+    client._set_spend_logs_row_count_in_proxy_state = AsyncMock()
+    client.start_db_health_watchdog_task = AsyncMock()
+    client.health_check = AsyncMock(side_effect=health_check_error)
+    return client
+
+
+async def _run_setup_prisma_client(mock_client):
+    from litellm.proxy.proxy_server import ProxyStartupEvent
+
+    with patch.object(proxy_server_module, "PrismaClient", return_value=mock_client):
+        result = await ProxyStartupEvent._setup_prisma_client(
+            database_url="postgresql://litellm:litellm@localhost:5432/litellm",
+            proxy_logging_obj=MagicMock(),
+            user_api_key_cache=DualCache(),
+        )
+    await asyncio.sleep(0.05)
+    return result
+
+
+@pytest.mark.asyncio
+async def test_setup_prisma_client_retains_connected_client_when_startup_health_check_fails(
+    monkeypatch,
+):
+    """A transient failure of the startup ``SELECT 1`` must not discard a client
+    whose ``connect()`` already succeeded.
+
+    Discarding it assigns ``None`` to the module-level ``prisma_client`` for the
+    life of the process, so a database that came back a second later is never
+    used again until the proxy is restarted."""
+    monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False")
+    monkeypatch.setattr(
+        proxy_server_module,
+        "general_settings",
+        {"allow_requests_on_db_unavailable": True},
+    )
+
+    mock_client = _mock_startup_prisma_client(
+        health_check_error=httpx.ReadTimeout("startup health check timed out")
+    )
+    result = await _run_setup_prisma_client(mock_client)
+
+    assert mock_client.connect.await_count == 1
+    assert mock_client.health_check.await_count == 1
+    assert result is mock_client
+
+
+@pytest.mark.asyncio
+async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_check(
+    monkeypatch,
+):
+    """The health watchdog is the only thing that reconnects a dropped DB, so it
+    has to be armed before the startup health check can fail.
+
+    Armed after, the single failure it exists to recover from is exactly the one
+    that skips it, and recovery never happens."""
+    monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False")
+    monkeypatch.setattr(
+        proxy_server_module,
+        "general_settings",
+        {"allow_requests_on_db_unavailable": True},
+    )
+
+    mock_client = _mock_startup_prisma_client(
+        health_check_error=httpx.ReadTimeout("startup health check timed out")
+    )
+    call_order = MagicMock()
+    call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog")
+    call_order.attach_mock(mock_client.health_check, "health_check")
+
+    await _run_setup_prisma_client(mock_client)
+
+    assert mock_client.start_db_health_watchdog_task.await_count == 1
+    assert [call[0] for call in call_order.mock_calls] == ["watchdog", "health_check"]
+
+
+@pytest.mark.asyncio
+async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(monkeypatch):
+    """Without ``allow_requests_on_db_unavailable`` a failed startup health check
+    must still hard-fail startup. Retaining the client is a fallback for
+    operators who opted into serving traffic without a database, never a way to
+    boot a proxy whose DB never answered."""
+    monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False")
+    monkeypatch.setattr(
+        proxy_server_module,
+        "general_settings",
+        {"allow_requests_on_db_unavailable": False},
+    )
+
+    mock_client = _mock_startup_prisma_client(
+        health_check_error=httpx.ReadTimeout("startup health check timed out")
+    )
+    with pytest.raises(httpx.ReadTimeout):
+        await _run_setup_prisma_client(mock_client)
+
+
+@pytest.mark.asyncio
+async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkeypatch):
+    """Retaining only ever applies to a client that connected. If ``connect()``
+    failed there is no usable client and no watchdog to recover it, so the caller
+    must still get ``None``."""
+    monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False")
+    monkeypatch.setattr(
+        proxy_server_module,
+        "general_settings",
+        {"allow_requests_on_db_unavailable": True},
+    )
+
+    mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused"))
+    result = await _run_setup_prisma_client(mock_client)
+
+    assert result is None
+    assert mock_client.start_db_health_watchdog_task.await_count == 0
+    assert mock_client.health_check.await_count == 0

tests/test_litellm/proxy/test_proxy_utils.py

@@ -1085,3 +1085,79 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks):
             request_data={"mcp_tool_name": "echo"},
             user_api_key_dict=None,
         )
+
+
+@pytest.mark.asyncio
+async def test_prisma_health_check_failure_names_itself_at_operator_visible_level(caplog):
+    """A failing DB health check has to name the check that failed, at a level
+    operators actually run at.
+
+    Reporting it as ``disconnect()`` sends anyone grepping the logs to the wrong
+    function and reads as "the check never ran", and reporting it only at debug
+    level hides a database fault behind a flag nobody enables in production."""
+    import logging
+    from unittest.mock import AsyncMock
+
+    from litellm.proxy.utils import PrismaClient
+
+    client = MagicMock()
+    client.db.query_raw = AsyncMock(side_effect=Exception("connection refused"))
+    client.proxy_logging_obj.failure_handler = AsyncMock()
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        with pytest.raises(Exception, match="connection refused"):
+            await PrismaClient.health_check(client)
+
+    assert "health_check()" in caplog.text
+    assert "disconnect()" not in caplog.text
+    assert "connection refused" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_prisma_connect_failure_is_reported_at_operator_visible_level(caplog):
+    """The sibling connect failure is labelled correctly but was equally
+    invisible. A database the proxy could not connect to at startup must not be
+    a debug-only record."""
+    import logging
+    from unittest.mock import AsyncMock
+
+    from litellm.proxy.utils import PrismaClient
+
+    client = MagicMock()
+    client.db.is_connected = MagicMock(return_value=False)
+    client.db.connect = AsyncMock(side_effect=Exception("could not reach database"))
+    client.proxy_logging_obj.failure_handler = AsyncMock()
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        with pytest.raises(Exception, match="could not reach database"):
+            await PrismaClient.connect(client)
+
+    assert "connect()" in caplog.text
+    assert "could not reach database" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_prisma_health_check_failure_redacts_database_credentials(caplog):
+    """Raising the level must not widen what reaches the logs. The exception
+    text can carry a full connection string, so the credential has to be gone
+    from the emitted record."""
+    import logging
+    from unittest.mock import AsyncMock
+
+    from litellm.proxy.utils import PrismaClient
+
+    client = MagicMock()
+    client.db.query_raw = AsyncMock(
+        side_effect=Exception("could not connect to postgresql://admin:hunter2@db.internal:5432/litellm")
+    )
+    client.proxy_logging_obj.failure_handler = AsyncMock()
+
+    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
+        with pytest.raises(Exception):
+            await PrismaClient.health_check(client)
+
+    emitted = [record.getMessage() for record in caplog.records if record.name == "LiteLLM Proxy"]
+
+    assert emitted
+    assert all("hunter2" not in message for message in emitted)
+    assert any("postgresql://REDACTED@db.internal" in message for message in emitted)