mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #41974 from BerriAI/litellm_fix_startup_view_creation_race
fix(proxy): wait for the spend-log table before creating startup views
This commit is contained in:
commit
9d7f77988a
7 changed files with 411 additions and 38 deletions
|
|
@ -1510,6 +1510,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()
|
||||
|
|
@ -10808,14 +10814,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,
|
||||
|
|
@ -269,6 +270,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: ...
|
||||
|
||||
|
|
@ -4297,6 +4307,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
|
||||
|
|
@ -6343,6 +6354,71 @@ 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()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._view_setup_task
|
||||
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", _VIEW_SETUP_GATE_TABLE
|
||||
)
|
||||
return "table_missing"
|
||||
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)
|
||||
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.",
|
||||
_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", _VIEW_SETUP_GATE_TABLE)
|
||||
)
|
||||
return rows[0]["present"]
|
||||
|
||||
async def _db_health_watchdog_loop(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -359,6 +359,17 @@ 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:
|
||||
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]}"
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -1810,6 +1810,37 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke
|
|||
assert announced == []
|
||||
|
||||
|
||||
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):
|
||||
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
|
||||
|
|
@ -13608,6 +13639,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
|
||||
|
|
@ -13669,13 +13701,35 @@ 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):
|
||||
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,221 @@ 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) -> None:
|
||||
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", "row_count", "views"],
|
||||
"probe_args": (_PROBE_SQL, '"LiteLLM_SpendLogs"'),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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")
|
||||
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_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
|
||||
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:
|
||||
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", "row_count", "views", "probe", "row_count", "views"],
|
||||
}
|
||||
|
||||
|
||||
@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", "row_count", "views"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout(
|
||||
prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
_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": '"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:
|
||||
_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