From 2863559ba8a422df87981e108c486316927ea859 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 08:36:22 -0700 Subject: [PATCH 1/3] fix(proxy): wait for the spend-log table before creating startup views On a fresh database where the migrations run in a separate job while the proxy boots with DISABLE_SCHEMA_UPDATE=true, the startup view check ran as a fire-and-forget task, used up its three 10 second retries before LiteLLM_SpendLogs existed, and died with an unretrieved exception. The spend views were never created, so the /global/spend routes returned 500 until the pod was restarted PrismaClient now holds a view setup task. It polls to_regclass for the spend-log table every 5 seconds, creates the views and loads the spend log row count once the table is there, keeps polling if an attempt raises while the schema is still settling, and logs an ERROR with the last failure if nothing worked after 15 minutes. Proxy shutdown cancels the task The spend route e2e tests for the five view-backed routes are no longer skipped and wait for the views through the harness convergence helper --- litellm/proxy/proxy_server.py | 15 +- litellm/proxy/utils.py | 83 ++++++ .../spend_tracking/spend_e2e_client.py | 13 + .../spend_tracking/test_spend_routes.py | 31 +-- tests/test_litellm/proxy/test_proxy_server.py | 63 ++++- .../test_prisma_client_lifecycle.py | 239 ++++++++++++++++++ 6 files changed, 418 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 36fdea605c2..a4d99d9859e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1464,6 +1464,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): + try: + await prisma_client.stop_view_setup_task() + except Exception as e: + verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) + await _drain_spend_event_producer_on_shutdown() await flush_spend_counters_on_shutdown() @@ -10635,14 +10641,7 @@ class ProxyStartupEvent: if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): await prisma_client.db.start_token_refresh_task() - ## 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 + prisma_client.start_view_setup_task() if hasattr(prisma_client, "start_db_health_watchdog_task"): await prisma_client.start_db_health_watchdog_task() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b078a65759e..7f5609b2e39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,6 +38,7 @@ from typing import ( Literal, Optional, Protocol, + TypeAlias, TypeVar, Union, cast, @@ -268,6 +269,15 @@ class _RelTuplesRow(TypedDict): reltuples: ReadOnly[int] +_VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 +_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 +_VIEW_SETUP_GATE_TABLE: Final = "LiteLLM_SpendLogs" +_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) + +_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] +_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception + + class _EndUserBatchTable(Protocol): def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... @@ -4284,6 +4294,7 @@ class PrismaClient: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: asyncio.Task | None = None + self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None self._db_last_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) self._db_read_only_recreate_ts: float = 0.0 @@ -6330,6 +6341,78 @@ class PrismaClient: self._db_health_watchdog_task = None verbose_proxy_logger.info("Stopped Prisma DB health watchdog") + def start_view_setup_task(self) -> None: + if self._view_setup_task is not None: + return + self._view_setup_task = asyncio.create_task(self._run_view_setup()) + + async def stop_view_setup_task(self) -> None: + if self._view_setup_task is None: + return + self._view_setup_task.cancel() + try: + await self._view_setup_task + except asyncio.CancelledError: + pass + self._view_setup_task = None + + async def _run_view_setup( + self, + poll_interval_seconds: float = _VIEW_SETUP_POLL_INTERVAL_SECONDS, + deadline_seconds: float = _VIEW_SETUP_DEADLINE_SECONDS, + ) -> _ViewSetupOutcome: + deadline: Final = time.monotonic() + deadline_seconds + while True: + if (attempt := await self._attempt_view_setup()) == "ready": + return "ready" + if time.monotonic() >= deadline: + self._log_view_setup_timeout(attempt, deadline_seconds) + return "timed_out" + await asyncio.sleep(poll_interval_seconds) + + async def _attempt_view_setup(self) -> _ViewSetupAttempt: + try: + if not await self._view_setup_gate_table_present(): + verbose_proxy_logger.debug( + "Waiting for table %s before creating the spend views", self._view_setup_gate_table() + ) + return "table_missing" + await self.check_view_exists() + await self._set_spend_logs_row_count_in_proxy_state() + return "ready" + except Exception as e: + verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e) + return e + + def _log_view_setup_timeout( + self, last_attempt: Literal["table_missing"] | Exception, deadline_seconds: float + ) -> None: + if isinstance(last_attempt, Exception): + verbose_proxy_logger.error( + "Gave up creating the spend views after %ss; the last attempt failed with: %s. " + "Fix that error and restart the proxy.", + deadline_seconds, + last_attempt, + ) + return + verbose_proxy_logger.error( + "Gave up creating the spend views: table %s did not appear within %ss. " + "Run the database migrations against this database and restart the proxy.", + self._view_setup_gate_table(), + deadline_seconds, + ) + + async def _view_setup_gate_table_present(self) -> bool: + rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( + await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", self._view_setup_gate_table()) + ) + return rows[0]["present"] + + @staticmethod + def _view_setup_gate_table() -> str: + pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") + return f'"{pg_schema}"."{_VIEW_SETUP_GATE_TABLE}"' + async def _db_health_watchdog_loop(self) -> None: while True: try: diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 9ac97f57f47..233aa81c2af 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -359,6 +359,19 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: + """Re-probe a route that depends on startup work the proxy finishes after it + starts serving, such as the spend views it creates once migrations land.""" + outcome: Final = await_converged( + lambda: self.probe(path, params=params), + converged=lambda result: result.healthy, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: return unwrap( self.proxy.transport.post( diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py index 8cb3e3927f0..67fd88bc84d 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py @@ -17,9 +17,11 @@ fast: no batch-write wait, no provider calls. """ from datetime import datetime, timedelta, timezone +from typing import Final import pytest +from e2e_http import ProbeResult from models import DateRangeParams from spend_e2e_client import SpendClient @@ -72,15 +74,10 @@ SPEND_ROUTES = ( _SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") -_MISSING_VIEW_SKIP = pytest.mark.skip( - reason=( - "LIT-5211: on a fresh database the proxy's startup view creation can lose the race " - "against schema migrations, leaving MonthlyGlobalSpend/DailyTagSpend/Last30d* views " - "missing and these routes 500ing until the views exist" - ) -) - -_VIEW_BACKED_ROUTES = frozenset( +# Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the +# proxy creates in the background once the schema migrations have landed, so on a +# fresh database they can 500 for a while after the proxy starts serving. +_VIEW_BACKED_ROUTES: Final = frozenset( ( "/global/spend", "/global/spend/keys", @@ -98,15 +95,15 @@ def _date_range() -> DateRangeParams: return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) -@pytest.mark.parametrize( - "route", - tuple( - pytest.param(route, marks=_MISSING_VIEW_SKIP) if route in _VIEW_BACKED_ROUTES else route - for route in SPEND_ROUTES - ), -) +def _probe(client: SpendClient, route: str) -> ProbeResult: + if route in _VIEW_BACKED_ROUTES: + return client.probe_until_healthy(route, params=_date_range()) + return client.probe(route, params=_date_range()) + + +@pytest.mark.parametrize("route", SPEND_ROUTES) def test_spend_route_responsive(client: SpendClient, route: str) -> None: - result = client.probe(route, params=_date_range()) + result = _probe(client, route) print(f"{route} -> {result.status_code}\n{result.body[:600]}") assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e2e2045e826..7ab142c5fd0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1628,6 +1628,40 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): assert master_key == test_resolved_key +class _ShutdownAwarePrisma(MockPrisma): + def __init__(self): + super().__init__() + self.stop_view_setup_task = AsyncMock() + + +@pytest.mark.asyncio +async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): + """The view setup task keeps polling for the spend-log table while migrations + run, so a shutdown inside that window has to cancel it rather than leave it + to die with the event loop.""" + import yaml + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + fake_prisma = _ShutdownAwarePrisma() + config_path = tmp_path / "config.yaml" + with open(config_path, "w") as f: + yaml.dump({"general_settings": {"master_key": "sk-12345"}}, f) + monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path)) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) + + async with proxy_startup_event(FastAPI()): + stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + + actual = { + "stopped_while_serving": stopped_while_serving, + "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + } + assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} + + def test_team_info_masking(): """ Test that sensitive team information is properly masked @@ -13307,6 +13341,7 @@ def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + client.start_view_setup_task = MagicMock() client.start_db_health_watchdog_task = AsyncMock() client.health_check = AsyncMock(side_effect=health_check_error) return client @@ -13368,13 +13403,39 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch 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_view_setup_task, "view_setup") 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"] + assert [call[0] for call in call_order.mock_calls] == ["view_setup", "watchdog", "health_check"] + + +@pytest.mark.asyncio +async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): + """View creation used to be two fire-and-forget ``asyncio.create_task`` calls + that raised and vanished when the migrations Job had not created + ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's + held task, which waits for the table, and must not call the two coroutines directly.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + + mock_client = _mock_startup_prisma_client() + result = await _run_setup_prisma_client(mock_client) + + actual = { + "result": result, + "view_setup_started": mock_client.start_view_setup_task.call_count, + "direct_view_creation": mock_client.check_view_exists.await_count, + "direct_row_count": mock_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "result": mock_client, + "view_setup_started": 1, + "direct_view_creation": 0, + "direct_row_count": 0, + } @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 18b02ac7772..bb34486771c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -5,11 +5,15 @@ Symbols pinned here: - ``PrismaClient.writer_db`` - ``PrismaClient.connect`` - ``PrismaClient.disconnect`` + - ``PrismaClient.start_view_setup_task`` + - ``PrismaClient.stop_view_setup_task`` + - ``PrismaClient._run_view_setup`` """ from __future__ import annotations import asyncio +import logging from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -17,6 +21,27 @@ import pytest from litellm.proxy.utils import PrismaClient +_PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" + + +def _absent() -> list[dict[str, bool]]: + return [{"present": False}] + + +def _present() -> list[dict[str, bool]]: + return [{"present": True}] + + +def _wire_view_setup(prisma_client: PrismaClient, probe: AsyncMock) -> MagicMock: + prisma_client.db.query_raw = probe + prisma_client.check_view_exists = AsyncMock() + prisma_client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + call_order = MagicMock() + call_order.attach_mock(probe, "probe") + call_order.attach_mock(prisma_client.check_view_exists, "views") + call_order.attach_mock(prisma_client._set_spend_logs_row_count_in_proxy_state, "row_count") + return call_order + @pytest.mark.asyncio async def test_prismaclient_init_wires_default_config( @@ -205,3 +230,217 @@ async def test_disconnect_raises_when_underlying_fails( prisma_client.db.disconnect = AsyncMock(side_effect=RuntimeError("disconnect boom")) with pytest.raises(RuntimeError, match="disconnect boom"): await prisma_client.disconnect() + + +@pytest.mark.asyncio +async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """On a fresh database the migrations Job can still be running when the proxy + boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the + table exists raised inside a fire-and-forget task and the views never appeared.""" + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + "probe_args": probe.await_args.args, + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "probe", "views", "row_count"], + "probe_args": (_PROBE_SQL, '"public"."LiteLLM_SpendLogs"'), + } + + +@pytest.mark.asyncio +async def test_view_setup_probes_the_configured_database_schema( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant") + probe = AsyncMock(return_value=_present()) + _wire_view_setup(prisma_client, probe) + + await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + assert probe.await_args.args == (_PROBE_SQL, '"litellm_tenant"."LiteLLM_SpendLogs"') + + +@pytest.mark.asyncio +async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "kept_polling": probe.await_count > 1, + "views_attempted": prisma_client.check_view_exists.await_count, + "row_count_attempted": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "outcome": "timed_out", + "kept_polling": True, + "views_attempted": 0, + "row_count_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: + """``LiteLLM_SpendLogs`` lands early in the migration set while + ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, + so the first attempt after the table appears can still fail.""" + probe = AsyncMock(return_value=_present()) + call_order = _wire_view_setup(prisma_client, probe) + prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None] + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "views", "probe", "views", "row_count"], + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_client: PrismaClient) -> None: + probe = AsyncMock(side_effect=[RuntimeError("connection reset"), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "views", "row_count"], + } + + +@pytest.mark.asyncio +async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_table": '"public"."LiteLLM_SpendLogs"' in errors[0], + "tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_table": True, + "tells_operator_to_migrate": True, + } + + +@pytest.mark.asyncio +async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + """A database role without CREATE on the schema fails every attempt even though + the table is there, so the timeout must blame that error, not missing migrations.""" + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_the_error": "permission denied for schema public" in errors[0], + "blames_missing_migrations": "did not appear" in errors[0], + "tells_operator_to_restart": "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_the_error": True, + "blames_missing_migrations": False, + "tells_operator_to_restart": True, + } + + +@pytest.mark.asyncio +async def test_run_view_setup_stays_quiet_when_views_are_ready( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + actual = { + "outcome": outcome, + "errors": [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR], + } + assert actual == {"outcome": "ready", "errors": []} + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_cancels_a_task_parked_between_polls(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + prisma_client.start_view_setup_task() + task = prisma_client._view_setup_task + await asyncio.sleep(0) + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "probed_before_parking": probe.await_count, + "task_cancelled": task is not None and task.cancelled(), + "reference_cleared": prisma_client._view_setup_task, + "views_attempted": prisma_client.check_view_exists.await_count, + } + assert actual == { + "probed_before_parking": 1, + "task_cancelled": True, + "reference_cleared": None, + "views_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_is_a_noop_without_a_task(prisma_client: PrismaClient) -> None: + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + assert prisma_client._view_setup_task is None + + +@pytest.mark.asyncio +async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: PrismaClient) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + prisma_client.start_view_setup_task() + first = prisma_client._view_setup_task + prisma_client.start_view_setup_task() + second = prisma_client._view_setup_task + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "first_is_task": isinstance(first, asyncio.Task), + "second_is_first": second is first, + } + assert actual == {"first_is_task": True, "second_is_first": True} From 1b8f704035a358d31962ca24d8ceda77d3dde935 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:17:18 -0700 Subject: [PATCH 2/3] fix(proxy): await the cancelled view setup task quietly and assert it starts at boot Use contextlib.suppress for the cancelled task in stop_view_setup_task, make the legacy prisma setup test inject a plain mock for the synchronous start_view_setup_task and assert it is called, and drop the docstrings the branch added to tests --- litellm/proxy/utils.py | 4 +--- .../spend_tracking/spend_e2e_client.py | 2 -- tests/proxy_unit_tests/test_proxy_server.py | 17 +++++------------ tests/test_litellm/proxy/test_proxy_server.py | 7 ------- .../test_prisma_client_lifecycle.py | 8 -------- 5 files changed, 6 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7f5609b2e39..37c5c8acec8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6350,10 +6350,8 @@ class PrismaClient: if self._view_setup_task is None: return self._view_setup_task.cancel() - try: + with contextlib.suppress(asyncio.CancelledError): await self._view_setup_task - except asyncio.CancelledError: - pass self._view_setup_task = None async def _run_view_setup( diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 233aa81c2af..b7f59fe5f89 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -360,8 +360,6 @@ class SpendClient: return self.proxy.transport.probe(path, params=params) def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: - """Re-probe a route that depends on startup work the proxy finishes after it - starts serving, such as the spend views it creates once migrations land.""" outcome: Final = await_converged( lambda: self.probe(path, params=params), converged=lambda result: result.healthy, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 47792b90b08..ed0380058a5 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2382,7 +2382,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py @pytest.mark.asyncio async def test_proxy_server_prisma_setup(): - from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_state + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging from litellm.caching import DualCache @@ -2393,35 +2393,28 @@ async def test_proxy_server_prisma_setup(): ) as mock_prisma_client: mock_client = mock_prisma_client.return_value # This is the mocked instance mock_client.connect = AsyncMock() # Mock the connect method - mock_client.check_view_exists = AsyncMock() # Mock the check_view_exists method + mock_client.start_view_setup_task = MagicMock() mock_client.health_check = AsyncMock() # Mock the health_check method - mock_client._set_spend_logs_row_count_in_proxy_state = ( - AsyncMock() - ) # Mock the _set_spend_logs_row_count_in_proxy_state method mock_client.start_db_health_watchdog_task = AsyncMock() # Mock the db attribute with start_token_refresh_task for RDS IAM token refresh mock_db = MagicMock() mock_db.start_token_refresh_task = AsyncMock() mock_client.db = mock_db - await ProxyStartupEvent._setup_prisma_client( + prisma_client = await ProxyStartupEvent._setup_prisma_client( database_url=os.getenv("DATABASE_URL"), proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) - # Verify our mocked methods were called + assert prisma_client is mock_client mock_client.connect.assert_called_once() - mock_client.check_view_exists.assert_called_once() + mock_client.start_view_setup_task.assert_called_once() # Note: This is REALLY IMPORTANT to check that the health check is called # This is how we ensure the DB is ready before proceeding mock_client.health_check.assert_called_once() - # check that the spend logs row count is set in proxy state - mock_client._set_spend_logs_row_count_in_proxy_state.assert_called_once() - assert proxy_state.get_proxy_state_variable("spend_logs_row_count") is not None - @pytest.mark.asyncio async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7ab142c5fd0..67f386ee457 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1636,9 +1636,6 @@ class _ShutdownAwarePrisma(MockPrisma): @pytest.mark.asyncio async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): - """The view setup task keeps polling for the spend-log table while migrations - run, so a shutdown inside that window has to cancel it rather than leave it - to die with the event loop.""" import yaml from fastapi import FastAPI @@ -13415,10 +13412,6 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch @pytest.mark.asyncio async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): - """View creation used to be two fire-and-forget ``asyncio.create_task`` calls - that raised and vanished when the migrations Job had not created - ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's - held task, which waits for the table, and must not call the two coroutines directly.""" monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") mock_client = _mock_startup_prisma_client() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index bb34486771c..c31713d5802 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -236,9 +236,6 @@ async def test_disconnect_raises_when_underlying_fails( async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: - """On a fresh database the migrations Job can still be running when the proxy - boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the - table exists raised inside a fire-and-forget task and the views never appeared.""" monkeypatch.delenv("DATABASE_SCHEMA", raising=False) probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) call_order = _wire_view_setup(prisma_client, probe) @@ -293,9 +290,6 @@ async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: P @pytest.mark.asyncio async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: - """``LiteLLM_SpendLogs`` lands early in the migration set while - ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, - so the first attempt after the table appears can still fail.""" probe = AsyncMock(return_value=_present()) call_order = _wire_view_setup(prisma_client, probe) prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None] @@ -358,8 +352,6 @@ async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table( prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture ) -> None: - """A database role without CREATE on the schema fails every attempt even though - the table is there, so the timeout must blame that error, not missing migrations.""" _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") From a114af2a26ed20f2a0dbb3fc452bc30f7d078d46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:54:20 -0700 Subject: [PATCH 3/3] fix(proxy): resolve the view setup gate through the search_path and set the row count before the views --- litellm/proxy/utils.py | 15 +++----- .../test_prisma_client_lifecycle.py | 38 ++++++++++++------- 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 37c5c8acec8..b45bcc3bdc3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -271,7 +271,7 @@ class _RelTuplesRow(TypedDict): _VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 -_VIEW_SETUP_GATE_TABLE: Final = "LiteLLM_SpendLogs" +_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] @@ -6372,11 +6372,11 @@ class PrismaClient: try: if not await self._view_setup_gate_table_present(): verbose_proxy_logger.debug( - "Waiting for table %s before creating the spend views", self._view_setup_gate_table() + "Waiting for table %s before creating the spend views", _VIEW_SETUP_GATE_TABLE ) return "table_missing" - await self.check_view_exists() await self._set_spend_logs_row_count_in_proxy_state() + await self.check_view_exists() return "ready" except Exception as e: verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e) @@ -6396,21 +6396,16 @@ class PrismaClient: verbose_proxy_logger.error( "Gave up creating the spend views: table %s did not appear within %ss. " "Run the database migrations against this database and restart the proxy.", - self._view_setup_gate_table(), + _VIEW_SETUP_GATE_TABLE, deadline_seconds, ) async def _view_setup_gate_table_present(self) -> bool: rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( - await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", self._view_setup_gate_table()) + await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) ) return rows[0]["present"] - @staticmethod - def _view_setup_gate_table() -> str: - pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public") - return f'"{pg_schema}"."{_VIEW_SETUP_GATE_TABLE}"' - async def _db_health_watchdog_loop(self) -> None: while True: try: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index c31713d5802..9aa57c7a19b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -233,10 +233,7 @@ async def test_disconnect_raises_when_underlying_fails( @pytest.mark.asyncio -async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( - prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("DATABASE_SCHEMA", raising=False) +async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views(prisma_client: PrismaClient) -> None: probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) call_order = _wire_view_setup(prisma_client, probe) @@ -249,13 +246,13 @@ async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( } assert actual == { "outcome": "ready", - "calls": ["probe", "probe", "probe", "views", "row_count"], - "probe_args": (_PROBE_SQL, '"public"."LiteLLM_SpendLogs"'), + "calls": ["probe", "probe", "probe", "row_count", "views"], + "probe_args": (_PROBE_SQL, '"LiteLLM_SpendLogs"'), } @pytest.mark.asyncio -async def test_view_setup_probes_the_configured_database_schema( +async def test_view_setup_probe_resolves_through_the_connection_search_path( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant") @@ -264,7 +261,23 @@ async def test_view_setup_probes_the_configured_database_schema( await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) - assert probe.await_args.args == (_PROBE_SQL, '"litellm_tenant"."LiteLLM_SpendLogs"') + assert probe.await_args.args == (_PROBE_SQL, '"LiteLLM_SpendLogs"') + + +@pytest.mark.asyncio +async def test_view_setup_sets_the_row_count_even_when_view_creation_keeps_failing( + prisma_client: PrismaClient, +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "row_count_set": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count >= 1, + } + assert actual == {"outcome": "timed_out", "row_count_set": True} @pytest.mark.asyncio @@ -302,7 +315,7 @@ async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_ } assert actual == { "outcome": "ready", - "calls": ["probe", "views", "probe", "views", "row_count"], + "calls": ["probe", "row_count", "views", "probe", "row_count", "views"], } @@ -319,15 +332,14 @@ async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_clien } assert actual == { "outcome": "ready", - "calls": ["probe", "probe", "views", "row_count"], + "calls": ["probe", "probe", "row_count", "views"], } @pytest.mark.asyncio async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( - prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture ) -> None: - monkeypatch.delenv("DATABASE_SCHEMA", raising=False) _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): @@ -337,7 +349,7 @@ async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( actual = { "outcome": outcome, "error_count": len(errors), - "names_table": '"public"."LiteLLM_SpendLogs"' in errors[0], + "names_table": '"LiteLLM_SpendLogs"' in errors[0], "tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0], } assert actual == {