fix(proxy): recover from cached-plan errors by reconnecting the Prisma client (#29983)

(cherry picked from commit 3bd3951e37)

litellm/proxy/utils.py

@@ -2841,40 +2841,49 @@ class PrismaClient:
         self, sql_query: str, *args
     ) -> Optional[dict]:
         """
-        Execute a query with automatic fallback for PostgreSQL cached plan errors.
+        Execute a query, recovering once from PostgreSQL's "cached plan must not
+        change result type" error.
 
-        This handles the "cached plan must not change result type" error that occurs
-        during rolling deployments when schema changes are applied while old pods
-        still have cached query plans expecting the old schema.
+        That error surfaces during rolling deployments when a schema change
+        invalidates the prepared-statement plans that pooled connections still
+        hold. Clearing only the server-side plans with DEALLOCATE ALL makes
+        things worse: Prisma's query engine keeps a per-connection client-side
+        cache of prepared-statement names, so once the server drops a plan the
+        engine re-sends a name PostgreSQL no longer recognizes and the
+        connection breaks with `prepared statement "sN" does not exist`. With a
+        small pool that connection stays poisoned and every auth lookup fails.
 
-        Args:
-            sql_query: SQL query string to execute
+        Recreating the Prisma client kills the engine subprocess and drops the
+        server-side plans and the engine's client-side name cache together, so
+        the retried query is prepared fresh. We reconnect through
+        `attempt_db_reconnect`, which is singleflight: when a schema change
+        poisons every pooled connection at once, the first cached-plan error
+        recreates the client and the concurrent waiters reuse that single
+        recreate instead of racing to kill each other's fresh engine. We then
+        retry the identical query exactly once.
 
-        Returns:
-            Query result or None
+        The retry reuses the original query byte-for-byte. Mutating the SQL
+        (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache,
+        forcing a fresh plan on every request and pegging the database CPU.
 
-        Raises:
-            Original exception if not a cached plan error
+        If the reconnect is skipped because a recent reconnect is still within
+        its cooldown, the retry runs against the same connection and may fail
+        again; the get_data backoff decorator re-runs the lookup and a later
+        attempt reconnects once the cooldown elapses.
         """
         try:
             return await self.db.query_first(sql_query, *args)
         except Exception as e:
-            error_str = str(e)
-            if "cached plan must not change result type" in error_str:
-                # Force PostgreSQL to re-plan by invalidating the cache
-                # Add a unique comment to make the query different
-                sql_query_retry = sql_query.replace(
-                    "SELECT",
-                    f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */",
-                )
-                verbose_proxy_logger.warning(
-                    "PostgreSQL cached plan error detected for token lookup, "
-                    "retrying with fresh plan. This may occur during rolling deployments "
-                    "when schema changes are applied."
-                )
-                return await self.db.query_first(sql_query_retry, *args)
-            else:
+            if "cached plan must not change result type" not in str(e):
                 raise
+            verbose_proxy_logger.warning(
+                "PostgreSQL cached plan error detected for token lookup; "
+                "recreating the database connection and retrying with the same "
+                "query. This may occur during rolling deployments when schema "
+                "changes are applied."
+            )
+            await self.attempt_db_reconnect(reason="postgres_cached_plan_error")
+            return await self.db.query_first(sql_query, *args)
 
     @backoff.on_exception(
         backoff.expo,

tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py

@@ -0,0 +1,180 @@
+"""Shared fixtures for tests/test_litellm/proxy/utils/prisma_and_spend/.
+
+All fixtures used by PR2 test files live here. Do NOT add fixtures inside
+individual test files; if a fixture is missing, add it here and update the
+Notion plan.
+
+The PrismaClient is exercised against a fully-mocked Prisma stack: the
+``prisma.Prisma`` constructor and the writer/reader wrappers are patched
+before PrismaClient.__init__ runs so the init code paths execute without
+needing a generated Prisma client or a real database.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from dataclasses import dataclass, field
+from email.message import EmailMessage
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterator, List, Optional
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[5]))
+
+
+_PRISMA_TABLES: List[str] = [
+    "litellm_verificationtoken",
+    "litellm_teamtable",
+    "litellm_usertable",
+    "litellm_endusertable",
+    "litellm_organizationtable",
+    "litellm_proxymodeltable",
+    "litellm_modeltable",
+    "litellm_budgettable",
+    "litellm_spendlogs",
+    "litellm_config",
+    "litellm_usernotifications",
+    "litellm_healthchecktable",
+    "litellm_dailyuserspend",
+    "litellm_dailyteamspend",
+    "litellm_dailytagspend",
+    "litellm_managed_object_table",
+    "litellm_credentialstable",
+    "litellm_mcpservertable",
+    "litellm_audit_log",
+    "litellm_invitationlink",
+    "litellm_session_token_table",
+    "litellm_passthrough_endpoint_table",
+    "litellm_cron_job",
+    "litellm_passthrough_logs",
+    "litellm_promptstable",
+    "litellm_guardrailstable",
+    "litellm_managed_files",
+    "litellm_mcpusercredentials",
+    "litellm_objectpermissiontable",
+    "litellm_organizationmembership",
+]
+
+
+def _make_table_mock() -> MagicMock:
+    table = MagicMock()
+    table.find_unique = AsyncMock(return_value=None)
+    table.find_many = AsyncMock(return_value=[])
+    table.find_first = AsyncMock(return_value=None)
+    table.create = AsyncMock()
+    table.create_many = AsyncMock()
+    table.update = AsyncMock()
+    table.update_many = AsyncMock()
+    table.upsert = AsyncMock()
+    table.delete = AsyncMock()
+    table.delete_many = AsyncMock()
+    table.count = AsyncMock(return_value=0)
+    table.group_by = AsyncMock(return_value=[])
+    table.aggregate = AsyncMock(return_value={})
+    return table
+
+
+
+@pytest.fixture
+def mock_prisma_client() -> MagicMock:
+    """Bare ``db`` mock with all common LiteLLM_* tables stubbed.
+
+    Override individual return values in a test::
+
+        mock_prisma_client.db.litellm_usertable.find_unique.return_value = user
+    """
+    client = MagicMock(name="MockPrismaClient")
+    client.db = MagicMock(name="MockPrismaDB")
+    client.connect = AsyncMock()
+    client.disconnect = AsyncMock()
+    client.health_check = AsyncMock(return_value=[{"?column?": 1}])
+    client.proxy_logging_obj = MagicMock()
+    client.proxy_logging_obj.failure_handler = AsyncMock()
+    client.spend_log_transactions = []
+    client._spend_log_transactions_lock = asyncio.Lock()
+    client.jsonify_object = lambda data: dict(data)
+    client.db.is_connected = MagicMock(return_value=False)
+    client.db.connect = AsyncMock()
+    client.db.disconnect = AsyncMock()
+    client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
+    client.db.execute_raw = AsyncMock()
+    client.db.tx = MagicMock()
+    client.db.batch_ = MagicMock()
+    for table_name in _PRISMA_TABLES:
+        setattr(client.db, table_name, _make_table_mock())
+    return client
+
+
+@pytest.fixture
+def patched_prisma_import(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]:
+    """Replace ``prisma.Prisma`` and ``PrismaWrapper`` so PrismaClient.__init__
+    runs without a generated client. Yields the fake Prisma instance.
+
+    ``prisma`` raises RuntimeError (not AttributeError) for the missing
+    ``Prisma`` attribute, so ``monkeypatch.setattr`` can't probe it; assign
+    directly and restore in teardown.
+    """
+    import prisma as _prisma_pkg
+    import litellm.proxy.utils as _utils_mod
+
+    fake_prisma = MagicMock(name="FakePrisma")
+    fake_prisma.is_connected = MagicMock(return_value=False)
+    fake_prisma.connect = AsyncMock()
+    fake_prisma.disconnect = AsyncMock()
+
+    fake_prisma_factory = MagicMock(name="FakePrismaFactory", return_value=fake_prisma)
+    had_prisma_attr = "Prisma" in _prisma_pkg.__dict__
+    previous_prisma_attr = _prisma_pkg.__dict__.get("Prisma")
+    _prisma_pkg.Prisma = fake_prisma_factory  # type: ignore[attr-defined]
+
+    fake_wrapper = MagicMock(name="FakePrismaWrapper")
+    fake_wrapper.is_connected = MagicMock(return_value=False)
+    fake_wrapper.connect = AsyncMock()
+    fake_wrapper.disconnect = AsyncMock()
+    fake_wrapper.query_raw = AsyncMock(return_value=[{"?column?": 1}])
+
+    def _fake_wrapper_ctor(*args: Any, **kwargs: Any) -> MagicMock:
+        return fake_wrapper
+
+    monkeypatch.setattr(_utils_mod, "PrismaWrapper", _fake_wrapper_ctor)
+    fake_prisma.__wrapper__ = fake_wrapper
+    try:
+        yield fake_prisma
+    finally:
+        if had_prisma_attr:
+            _prisma_pkg.Prisma = previous_prisma_attr  # type: ignore[attr-defined]
+        else:
+            try:
+                del _prisma_pkg.Prisma  # type: ignore[attr-defined]
+            except AttributeError:
+                pass
+
+
+@pytest.fixture
+def prisma_client(
+    patched_prisma_import: MagicMock,
+    mock_prisma_client: MagicMock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> Any:
+    """Wired ``PrismaClient`` whose ``db`` attribute is the table mock.
+
+    The init runs through the real code path (testing the constructor's
+    config-attribute setup) and is then snapped to the easier-to-assert
+    table mock for downstream behavior pinning.
+    """
+    monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
+    monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False)
+    from litellm.proxy.utils import PrismaClient
+
+    proxy_logging_obj = MagicMock(name="MockProxyLogging")
+    proxy_logging_obj.failure_handler = AsyncMock()
+    pc = PrismaClient(
+        database_url="postgresql://test:test@localhost:5432/test",
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    pc.db = mock_prisma_client.db
+    return pc
+

tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py

@@ -0,0 +1,154 @@
+"""Pin ``PrismaClient`` read-side data operations.
+
+Symbols pinned here:
+  - ``PrismaClient.hash_token``
+  - ``PrismaClient.jsonify_object``
+  - ``PrismaClient.jsonify_team_object``
+  - ``PrismaClient.check_view_exists``
+  - ``PrismaClient.get_request_status``
+  - ``PrismaClient.get_generic_data``
+  - ``PrismaClient._query_first_with_cached_plan_fallback``
+  - ``PrismaClient.get_data``
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.proxy.utils import PrismaClient
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_happy_returns_row(
+    prisma_client: PrismaClient,
+) -> None:
+    expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
+    prisma_client.db.query_first = AsyncMock(return_value=expected)
+    prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
+    result = await prisma_client._query_first_with_cached_plan_fallback(
+        "SELECT * FROM x WHERE token = $1", "abc"
+    )
+    actual = {
+        "result": result,
+        "call_count": prisma_client.db.query_first.await_count,
+        "args": prisma_client.db.query_first.await_args.args,
+        "matches": result == expected,
+    }
+    assert actual == {
+        "result": expected,
+        "call_count": 1,
+        "args": ("SELECT * FROM x WHERE token = $1", "abc"),
+        "matches": True,
+    }
+    prisma_client.attempt_db_reconnect.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query(
+    prisma_client: PrismaClient,
+) -> None:
+    original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1'
+    expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0}
+    manager = MagicMock()
+    query_first = AsyncMock(
+        side_effect=[
+            RuntimeError("cached plan must not change result type"),
+            expected,
+        ]
+    )
+    reconnect = AsyncMock(return_value=True)
+    manager.attach_mock(query_first, "query_first")
+    manager.attach_mock(reconnect, "attempt_db_reconnect")
+    prisma_client.db.query_first = query_first
+    prisma_client.attempt_db_reconnect = reconnect
+
+    result = await prisma_client._query_first_with_cached_plan_fallback(
+        original_query, "abc"
+    )
+
+    assert result == expected
+    assert query_first.await_count == 2
+    first_call, retry_call = query_first.await_args_list
+    assert retry_call.args == first_call.args == (original_query, "abc")
+    reconnect.assert_awaited_once()
+    assert reconnect.await_args.kwargs.get("force", False) is False
+    assert [name for name, *_ in manager.mock_calls] == [
+        "query_first",
+        "attempt_db_reconnect",
+        "query_first",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_never_deallocates(
+    prisma_client: PrismaClient,
+) -> None:
+    expected = {"token": "abc"}
+    prisma_client.db.query_first = AsyncMock(
+        side_effect=[
+            RuntimeError("cached plan must not change result type"),
+            expected,
+        ]
+    )
+    prisma_client.db.execute_raw = AsyncMock(return_value=0)
+    prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
+
+    await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
+
+    prisma_client.db.execute_raw.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails(
+    prisma_client: PrismaClient,
+) -> None:
+    plan_error = RuntimeError("cached plan must not change result type")
+    prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error])
+    prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
+
+    with pytest.raises(RuntimeError, match="cached plan must not change result type"):
+        await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
+
+    assert prisma_client.db.query_first.await_count == 2
+    prisma_client.attempt_db_reconnect.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false(
+    prisma_client: PrismaClient,
+) -> None:
+    expected = {"token": "abc"}
+    prisma_client.db.query_first = AsyncMock(
+        side_effect=[
+            RuntimeError("cached plan must not change result type"),
+            expected,
+        ]
+    )
+    prisma_client.attempt_db_reconnect = AsyncMock(return_value=False)
+
+    result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
+
+    assert result == expected
+    assert prisma_client.db.query_first.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors(
+    prisma_client: PrismaClient,
+) -> None:
+    prisma_client.db.query_first = AsyncMock(
+        side_effect=RuntimeError("totally unrelated")
+    )
+    prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
+    with pytest.raises(RuntimeError, match="totally unrelated"):
+        await prisma_client._query_first_with_cached_plan_fallback("SELECT 1")
+    assert prisma_client.db.query_first.await_count == 1
+    prisma_client.attempt_db_reconnect.assert_not_awaited()
+