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] 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 == {