mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
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
This commit is contained in:
mateo-berri
parent
b8d837b2ef
commit
2863559ba8
6 changed files with 418 additions and 26 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue