scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 312

refactor(proxy): give database shutdown one lifespan owner (#48506)

Browse source 
* refactor(proxy): give database shutdown one lifespan owner

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(proxy): log database cleanup errors and bound shutdown waits

Database.aclose now logs a failure in either cleanup step instead of raising
out of the lifespan exit, which matches how main handled stop errors and keeps
setup_prisma_client from raising out of cleanup after a failed connect

Shutdown waits for an in-flight reconnect for a fixed five seconds instead of
the reconnect timeout plus grace, and joins engine wait threads with a fixed
limit instead of indefinitely. The threads are daemon threads, so a thread
still alive after the limit cannot keep the process from exiting

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yujonglee 2026-10-30 15:57:53 -07:00 • committed by GitHub
parent c54ee63be9
commit 2d79d5dba7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 664 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.repositories.database import Database
from litellm.tracing import TraceReceiver
Span = _Span | Any
@ -133,6 +134,7 @@ else:
class ProxyLifespanState(TypedDict):
tracing_receiver: ReadOnly["TraceReceiver | None"]
database: ReadOnly["Database | None"]
class ReconcileOutcome(NamedTuple):

10
litellm/proxy/db/routing_prisma_wrapper.py
View file

@ -4,6 +4,7 @@ to a writer client. Used when DATABASE_URL_READ_REPLICA is configured;
otherwise PrismaClient uses the writer-only PrismaWrapper directly.
"""
import asyncio
import os
from collections.abc import Callable
from datetime import timedelta
@ -233,8 +234,13 @@ class RoutingPrismaWrapper:
await self._reader.start_token_refresh_task()
async def stop_token_refresh_task(self) -> None:
await self._writer.stop_token_refresh_task()
await self._reader.stop_token_refresh_task()
results: Final = await asyncio.gather(
self._writer.stop_token_refresh_task(),
self._reader.stop_token_refresh_task(),
return_exceptions=True,
)
if (error := next((result for result in results if isinstance(result, BaseException)), None)) is not None:
raise error
async def recreate_prisma_client(
self,

122
litellm/proxy/proxy_server.py
View file

@ -174,6 +174,7 @@ from litellm.proxy.common_utils.token_count_settings import (
)
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields
from litellm.repositories import database as database_runtime
from litellm.router_strategy.auto_router.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
@ -1350,13 +1351,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
if worker_heartbeat is not None and prisma_client:
await worker_heartbeat.deregister()
if prisma_client:
# Drain the SGR fold first: it lives in memory, so an un-drained interval
# is lost, and a write attempted after disconnect raises
# ClientNotConnectedError rather than persisting anything. Ordering this
# inside the same guard is what keeps the two from drifting apart.
await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client))
verbose_proxy_logger.debug("Disconnecting from Prisma")
await prisma_client.disconnect()
else:
await _drain_audit_tasks_on_shutdown()
@ -1479,23 +1474,38 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None":
_initialize_shared_aiohttp_session: Final = initialize_shared_aiohttp_session
async def _connect_to_count_stored_values() -> SupportsRawQueries:
async def _connect_to_count_stored_values(resources: AsyncExitStack) -> SupportsRawQueries:
client: Final = prisma_client or PrismaClient(
database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj
)
if prisma_client is None:
_own_database(resources, client)
await client.connect()
return client.writer_db
def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None:
if client is None:
return None
database: Final = database_runtime.Database(prisma=client)
resources.push_async_callback(database.aclose)
return database
def _forget_prisma_client() -> None:
global prisma_client # noqa: PLW0603 # routes still read this global until they move to DatabaseDependency
prisma_client = None # rebind-ok: a closed client must not survive into the next lifespan entry
@asynccontextmanager
async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
with GracefulShutdownManager.capture_sigterm():
async with _proxy_startup_event(app) as state:
async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state:
yield state
@asynccontextmanager
async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]:
global \
prisma_client, \
master_key, \
@ -1514,6 +1524,9 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
shared_aiohttp_session
import json
resources.callback(_forget_prisma_client)
existing_database: Final = _own_database(resources, prisma_client)
init_verbose_loggers()
running_loop: Final = asyncio.get_running_loop()
@ -1592,21 +1605,24 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
if decoded_worker_config is not None:
await initialize_from_worker_config(decoded_worker_config)
enforce_master_key_boot_verdict(
await with_stored_secrets_counted(
master_key_boot_verdict(
master_key=master_key,
environment_master_key=os.getenv(MASTER_KEY_ENV_VAR),
general_settings=general_settings,
config_file_path=user_config_file_path,
override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True,
salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None,
database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None,
async with AsyncExitStack() as boot_database:
enforce_master_key_boot_verdict(
await with_stored_secrets_counted(
master_key_boot_verdict(
master_key=master_key,
environment_master_key=os.getenv(MASTER_KEY_ENV_VAR),
general_settings=general_settings,
config_file_path=user_config_file_path,
override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True,
salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None,
database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None,
),
count_values_encrypted_with=partial(
count_values_encrypted_with_or_none, partial(_connect_to_count_stored_values, boot_database)
),
),
count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values),
),
announce=announce_on_stderr_at_exit,
)
announce=announce_on_stderr_at_exit,
)
# check if DATABASE_URL in environment - load from there
if prisma_client is None:
@ -1619,6 +1635,8 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
)
)
database: Final = existing_database or _own_database(resources, prisma_client)
await migrate_if_requested(
environ=os.environ,
master_key=master_key,
@ -1628,15 +1646,16 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
)
if prisma_client is not None:
startup_prisma: Final = prisma_client
async def _run_pw_migration():
async def _run_pw_migration() -> None:
try:
result: Final = await migrate_passwords_to_scrypt_async(prisma_client)
result: Final = await migrate_passwords_to_scrypt_async(startup_prisma)
verbose_proxy_logger.info("Password migration: %s", result)
except Exception as e:
verbose_proxy_logger.warning("Password migration skipped: %s", e)
asyncio.create_task(_run_pw_migration())
await resources.enter_async_context(database_runtime.background_task(_run_pw_migration()))
async def _run_agent_grant_id_migration() -> None:
from litellm.proxy.agent_endpoints.agent_registry import (
@ -1647,7 +1666,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
for attempt in range(3):
try:
result = await global_agent_registry.migrate_legacy_grant_ids(
table=object_permission_table(prisma_client)
table=object_permission_table(startup_prisma)
)
if result.rewritten:
verbose_proxy_logger.info(
@ -1667,7 +1686,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
if attempt < 2:
await asyncio.sleep(5)
asyncio.create_task(_run_agent_grant_id_migration())
await resources.enter_async_context(database_runtime.background_task(_run_agent_grant_id_migration()))
## A coordination_redis block saved from the admin UI lives in the database,
## which is only reachable once the prisma client exists. Apply it here, before
@ -1853,7 +1872,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
enabled=tracing_enabled,
settings=tracing_settings,
) as receiver:
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database}
from litellm.proxy.admin_mcp import admin_mcp_lifespan
telemetry_settings: Final = load_telemetry_settings()
@ -1907,30 +1926,6 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat
except Exception as e:
verbose_proxy_logger.error("Error closing shared aiohttp session: %s", e)
# Shutdown event - stop RDS IAM token refresh background task
if (
prisma_client is not None
and hasattr(prisma_client, "db")
and hasattr(prisma_client.db, "stop_token_refresh_task")
):
try:
await prisma_client.db.stop_token_refresh_task()
except Exception as e:
verbose_proxy_logger.error("Error stopping token refresh task: %s", e)
# Shutdown event - stop Prisma DB health watchdog task
if prisma_client is not None and hasattr(prisma_client, "stop_db_health_watchdog_task"):
try:
await prisma_client.stop_db_health_watchdog_task()
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()
if telemetry_start is not None:
telemetry_start.cancel()
@ -12124,17 +12119,20 @@ class ProxyStartupEvent:
proxy_logging_obj: ProxyLogging,
user_api_key_cache: UserApiKeyCache,
) -> PrismaClient | None:
"""
- Sets up prisma client
- Adds necessary views to proxy
"""
if database_url is None:
return None
client: Final = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj)
async with AsyncExitStack() as setup_resources:
setup_resources.push_async_callback(database_runtime.Database(prisma=client).aclose)
connected: Final = await cls._initialize_prisma_client(client)
if connected is not None:
setup_resources.pop_all()
return connected
@staticmethod
async def _initialize_prisma_client(prisma_client: PrismaClient) -> PrismaClient | None:
connected_client: PrismaClient | None = None
try:
if database_url is None:
return None
prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj)
try:
await prisma_client.connect()
except Exception as e:

73
litellm/proxy/utils.py
View file

@ -5197,6 +5197,8 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam
_WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only"
_WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]])
_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600
_SHUTDOWN_RECONNECT_WAIT_SECONDS: Final = 5.0
_SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS: Final = 5.0
class _ForcedRecreateDeclined(Exception):
@ -5281,6 +5283,7 @@ class PrismaClient:
database_url: str,
proxy_logging_obj: ProxyLogging,
http_client: "HttpConfig | None" = None,
engine_exit_wait_seconds: float = _SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS,
):
## init logging object
self.baseline_accounting_transactions: list[BaselineAccountingRecord] = []
@ -5382,6 +5385,9 @@ class PrismaClient:
else:
self.db = writer_wrapper # Client to connect to Prisma db
self._db_reconnect_lock = asyncio.Lock()
self._closing: bool = False
self._engine_exit_wait_seconds: Final = engine_exit_wait_seconds
self._engine_tasks: frozenset[asyncio.Task[object]] = frozenset()
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
@ -5423,6 +5429,7 @@ class PrismaClient:
self._watching_engine: bool = False
self._engine_confirmed_dead: bool = False
self._engine_wait_thread: threading.Thread | None = None
self._engine_wait_threads: tuple[threading.Thread, ...] = ()
verbose_proxy_logger.debug("Success - Created Prisma Client")
@property
@ -6631,6 +6638,14 @@ class PrismaClient:
start_time: Final = time.time()
try:
await self.db.disconnect()
if self._closing:
await asyncio.gather(
*(
asyncio.to_thread(thread.join, self._engine_exit_wait_seconds)
for thread in self._engine_wait_threads
)
)
self._engine_wait_threads = ()
except Exception as e:
import traceback
@ -6649,6 +6664,43 @@ class PrismaClient:
)
raise e
def _spawn_engine_task(self, work: Coroutine[object, object, object]) -> None:
if self._closing:
work.close()
return
task: Final = asyncio.create_task(work)
self._engine_tasks = self._engine_tasks | {task}
task.add_done_callback(self._engine_task_finished)
def _engine_task_finished(self, task: asyncio.Task[object]) -> None:
self._engine_tasks = self._engine_tasks - {task}
if not task.cancelled() and (error := task.exception()) is not None:
verbose_proxy_logger.error("Prisma engine background task failed: %s", error)
async def _stop_engine_tasks(self) -> None:
tasks: Final = tuple(self._engine_tasks)
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
await self._wait_for_inflight_reconnect()
async def _wait_for_inflight_reconnect(self) -> None:
waiter: Final = asyncio.ensure_future(self._db_reconnect_lock.acquire())
waiter.add_done_callback(lambda acquired: acquired.cancelled() or self._db_reconnect_lock.release())
done, _pending = await asyncio.wait({waiter}, timeout=_SHUTDOWN_RECONNECT_WAIT_SECONDS)
if waiter not in done:
waiter.cancel()
verbose_proxy_logger.warning("Shutdown stopped waiting for an in-flight DB reconnect")
async def stop_background_tasks(self) -> None:
self._closing = True
self._stop_engine_watcher()
async with contextlib.AsyncExitStack() as cleanup:
cleanup.push_async_callback(self._stop_engine_tasks)
cleanup.push_async_callback(self.db.stop_token_refresh_task)
cleanup.push_async_callback(self.stop_view_setup_task)
cleanup.push_async_callback(self.stop_db_health_watchdog_task)
def _get_engine_pid(self) -> int:
"""Get the PID of the writer's engine subprocess, or 0 if unavailable.
@ -6742,7 +6794,7 @@ class PrismaClient:
self._engine_confirmed_dead = True
self._reap_all_zombies()
self._cleanup_engine_watcher()
asyncio.create_task(
self._spawn_engine_task(
self.attempt_db_reconnect(
reason="engine_process_death",
force=True,
@ -6762,6 +6814,9 @@ class PrismaClient:
name=f"prisma-engine-waitpid-{pid}",
)
thread.start()
self._engine_wait_threads = tuple(thread for thread in self._engine_wait_threads if thread.is_alive()) + (
thread,
)
self._engine_wait_thread = thread
return True
@ -6823,7 +6878,7 @@ class PrismaClient:
self._engine_confirmed_dead = True
self._reap_all_zombies()
self._cleanup_engine_watcher()
asyncio.create_task(
self._spawn_engine_task(
self.attempt_db_reconnect(
reason="engine_process_death",
force=True,
@ -6885,7 +6940,7 @@ class PrismaClient:
self._engine_confirmed_dead = True
self._reap_all_zombies()
self._cleanup_engine_watcher()
asyncio.create_task(
self._spawn_engine_task(
self.attempt_db_reconnect(
reason="engine_process_death",
force=True,
@ -6958,7 +7013,7 @@ class PrismaClient:
waitpid thread nor pidfd are available.
"""
if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None:
if self._closing or self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None:
return
pid: Final = self._get_engine_pid()
if pid == 0:
@ -6985,7 +7040,7 @@ class PrismaClient:
pid,
)
self._watching_engine = True
asyncio.create_task(self._poll_engine_proc())
self._spawn_engine_task(self._poll_engine_proc())
def _stop_engine_watcher(self) -> None:
"""Stop watching the engine process and clean up all resources."""
@ -7007,7 +7062,7 @@ class PrismaClient:
"""
self._engine_confirmed_dead = False
self._cleanup_engine_watcher()
asyncio.create_task(self._start_engine_watcher())
self._spawn_engine_task(self._start_engine_watcher())
async def _run_reconnect_cycle(
self,
@ -7214,6 +7269,8 @@ class PrismaClient:
force_recreate: bool = False,
stale_read_engine: "_StaleReadEngine | None" = None,
) -> bool:
if self._closing:
return False
now: Final = time.time()
if (
force is False
@ -7404,7 +7461,7 @@ class PrismaClient:
- A periodic SELECT 1 probe that triggers reconnect on network/connection failure.
- A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.
"""
if self._db_health_watchdog_enabled is not True:
if self._closing or self._db_health_watchdog_enabled is not True:
verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED")
return
if self._db_health_watchdog_task is not None:
@ -7438,7 +7495,7 @@ class PrismaClient:
verbose_proxy_logger.info("Stopped Prisma DB health watchdog")
def start_view_setup_task(self) -> None:
if self._view_setup_task is not None:
if self._closing or self._view_setup_task is not None:
return
self._view_setup_task = asyncio.create_task(self._run_view_setup())

36
litellm/repositories/database.py Normal file
View file

@ -0,0 +1,36 @@
import asyncio
from collections.abc import AsyncGenerator, Coroutine
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
@dataclass(frozen=True, slots=True)
class Database:
prisma: "PrismaClient"
async def aclose(self) -> None:
verbose_proxy_logger.debug("Disconnecting from Prisma")
try:
await self.prisma.stop_background_tasks()
except Exception:
verbose_proxy_logger.exception("Stopping database background tasks failed")
try:
await self.prisma.disconnect()
except Exception:
verbose_proxy_logger.exception("Database disconnect failed")
@asynccontextmanager
async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]:
task: Final = asyncio.create_task(work)
try:
yield
finally:
task.cancel()
await asyncio.gather(task, return_exceptions=True)

158
tests/integration/database/test_proxy_database_lifespan.py Normal file
View file

@ -0,0 +1,158 @@
import asyncio
import os
import subprocess
import sys
import threading
import uuid
from pathlib import Path
from typing import Final, Literal, TypeAlias
import psutil
import pytest
from tests.integration._support.client import eventually
from tests.integration._support.database import read_rows, scratch_database
pytestmark: Final = pytest.mark.timeout(180)
ROOT: Final = Path(__file__).resolve().parents[3]
PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234"
Mode: TypeAlias = Literal["late", "unsafe"]
CONFIG: Final = (
"model_list: []\n"
"general_settings:\n"
" disable_prisma_schema_update: true\n"
" disable_spend_logs: true\n"
" disable_model_info_refresh: true\n"
)
def _assert_database_idle() -> None:
assert (
eventually(
lambda: read_rows(
"SELECT pid FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()",
(),
),
lambda rows: not rows,
seconds=5,
)
== []
)
assert psutil.Process().children(recursive=True) == []
assert tuple(thread.name for thread in threading.enumerate() if thread.name.startswith("prisma-engine-")) == ()
async def _assert_lifespan_closed() -> None:
from litellm.proxy import proxy_server
await asyncio.sleep(0)
assert proxy_server.prisma_client is None
assert (
tuple(
task.get_coro().__qualname__
for task in asyncio.all_tasks()
if task.get_coro().__qualname__.startswith(
("PrismaClient.", "PrismaWrapper.", "_proxy_lifespan.<locals>._run_")
)
)
== ()
)
await asyncio.to_thread(_assert_database_idle)
async def _exercise_lifespans(mode: Mode) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError
expected: Final = UnsafeMasterKeyError if mode == "unsafe" else ValueError
match: Final = "master key" if mode == "unsafe" else "use_redis_transaction_buffer"
with pytest.raises(expected, match=match):
async with proxy_server.proxy_startup_event(proxy_server.app):
pytest.fail("Invalid configuration started the proxy")
await _assert_lifespan_closed()
Path(os.environ["CONFIG_FILE_PATH"]).write_text(CONFIG)
with pytest.MonkeyPatch.context() as environment:
environment.setenv("LITELLM_MASTER_KEY", f"sk-{uuid.uuid4().hex}")
for cancel in (False, True, False):
try:
async with proxy_server.proxy_startup_event(proxy_server.app) as state:
database: Final = state["database"]
assert database is not None
assert await database.prisma.writer_db.query_raw("SELECT 42 AS answer") == [{"answer": 42}]
assert await database.prisma.read_db.query_raw("SELECT 43 AS answer") == [{"answer": 43}]
assert psutil.Process().children(recursive=True)
if cancel:
raise asyncio.CancelledError
except asyncio.CancelledError:
if not cancel:
raise
await _assert_lifespan_closed()
@pytest.mark.parametrize("mode", ("late", "unsafe"))
@pytest.mark.parametrize("replica", (False, True), ids=("writer", "writer-and-reader"))
def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, mode: Mode, replica: bool) -> None:
config: Final = tmp_path / "config.yaml"
config.write_text(CONFIG + (" use_redis_transaction_buffer: true\n" if mode == "late" else ""))
inherited: Final = {
key: value
for key, value in os.environ.items()
if not key.startswith("REDIS_")
and key
not in (
"WORKER_CONFIG",
"DATABASE_URL_READ_REPLICA",
"LITELLM_SALT_KEY",
"LITELLM_MIGRATE_FROM_MASTER_KEY",
"LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY",
"INTEGRATION_UPSTREAM_URL",
)
}
with scratch_database("lifespan") as database_url:
environment: Final = {
**inherited,
"DATABASE_URL": database_url,
"DIRECT_URL": database_url,
"CONFIG_FILE_PATH": str(config),
"LITELLM_MASTER_KEY": PUBLICLY_KNOWN_KEY if mode == "unsafe" else f"sk-{uuid.uuid4().hex}",
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
"LITELLM_TELEMETRY_ENABLED": "false",
"PRISMA_HEALTH_WATCHDOG_ENABLED": "true",
"PYTHONPATH": os.pathsep.join((str(ROOT), str(ROOT / "enterprise"))),
**({"DATABASE_URL_READ_REPLICA": database_url} if replica else {}),
}
schema: Final = subprocess.run(
(
sys.executable,
"-P",
"-m",
"prisma",
"db",
"push",
"--skip-generate",
"--schema",
str(ROOT / "litellm/proxy/schema.prisma"),
),
cwd=ROOT,
env=environment,
capture_output=True,
text=True,
timeout=60,
)
assert schema.returncode == 0, schema.stderr
probe: Final = subprocess.run(
(
sys.executable,
"-P",
"-c",
"import asyncio; from tests.integration.database.test_proxy_database_lifespan import _exercise_lifespans; "
f"asyncio.run(_exercise_lifespans({mode!r}))",
),
cwd=ROOT,
env=environment,
capture_output=True,
text=True,
timeout=90,
)
assert probe.returncode == 0, probe.stdout + probe.stderr

49
tests/unit/proxy/db/test_routing_prisma_wrapper.py
View file

@ -254,6 +254,55 @@ def test_token_refresh_delegates_to_both_writer_and_reader():
reader.stop_token_refresh_task.assert_awaited_once()
def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None:
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
writer: Final = MagicMock()
reader: Final = MagicMock()
writer.stop_token_refresh_task = AsyncMock(side_effect=RuntimeError("refresh failed"))
reader.stop_token_refresh_task = AsyncMock()
routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader)
with pytest.raises(RuntimeError, match="refresh failed"):
asyncio.run(routing.stop_token_refresh_task())
reader.stop_token_refresh_task.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("error", (RuntimeError("refresh failed"), asyncio.CancelledError()))
async def test_reader_cleanup_finishes_before_writer_stop_failure_propagates(error: BaseException) -> None:
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
reader_stopping: Final = asyncio.Event()
release_reader: Final = asyncio.Event()
reader_stopped: Final = asyncio.Event()
async def stop_reader() -> None:
reader_stopping.set()
await release_reader.wait()
reader_stopped.set()
writer: Final = MagicMock()
reader: Final = MagicMock()
writer.stop_token_refresh_task = AsyncMock(side_effect=error)
reader.stop_token_refresh_task = AsyncMock(side_effect=stop_reader)
routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader)
stopping: Final = asyncio.create_task(routing.stop_token_refresh_task())
try:
await reader_stopping.wait()
await asyncio.sleep(0)
assert not stopping.done()
release_reader.set()
with pytest.raises(type(error)):
await stopping
assert reader_stopped.is_set()
finally:
release_reader.set()
stopping.cancel()
await asyncio.gather(stopping, return_exceptions=True)
def test_routed_actions_falls_back_to_writer_for_unknown_methods():
from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions

37
tests/unit/proxy/proxy_server/test_lifecycle.py
View file

@ -24,6 +24,7 @@ import logging
import os
import subprocess
from collections.abc import Awaitable, Callable
from contextlib import aclosing
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -35,7 +36,6 @@ from typing_extensions import TypedDict
import litellm.proxy.proxy_server as ps
from litellm.proxy.management_helpers.audit_logs import track_audit_task
from litellm.proxy.proxy_server import (
_proxy_startup_event,
ProxyStartupEvent,
_resolve_pydantic_type,
_resolve_typed_dict_type,
@ -51,6 +51,7 @@ from litellm.proxy.proxy_server import (
run_ollama_serve,
save_worker_config,
)
from litellm.repositories import database as database_runtime
from .conftest import normalize
@ -101,9 +102,10 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises():
@pytest.mark.asyncio
async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch):
fake_prisma = MagicMock()
fake_prisma.disconnect = AsyncMock()
fake_prisma.stop_background_tasks = AsyncMock()
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
@ -117,7 +119,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
monkeypatch.setattr(litellm, "cache", None, raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
await proxy_shutdown_event()
async with aclosing(database_runtime.Database(prisma=fake_prisma)):
await proxy_shutdown_event()
fake_prisma.disconnect.assert_not_awaited()
observed = {
"disconnect_called": fake_prisma.disconnect.await_count == 1,
@ -204,6 +208,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
fake_prisma = MagicMock()
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
fake_prisma.stop_background_tasks = AsyncMock()
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
async def _record_flush(client, accumulator):
@ -222,7 +227,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
monkeypatch.setattr(litellm, "cache", None, raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
await proxy_shutdown_event()
async with aclosing(database_runtime.Database(prisma=fake_prisma)):
await proxy_shutdown_event()
assert calls == ["flush", "disconnect"]
@ -239,6 +245,7 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor
fake_prisma = MagicMock()
fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
fake_prisma.stop_background_tasks = AsyncMock()
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
async def _record_flush(client, accumulator):
@ -262,7 +269,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts()))
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
async with aclosing(database_runtime.Database(prisma=fake_prisma)):
await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
assert calls == ["flush", "audit", "disconnect"]
@ -318,9 +326,10 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t
@pytest.mark.asyncio
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
async def test_database_lifespan_logs_disconnect_error_after_shutdown_instead_of_raising(monkeypatch, caplog):
fake_prisma = MagicMock()
fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
fake_prisma.stop_background_tasks = AsyncMock()
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
fake_jwt = MagicMock()
@ -332,8 +341,12 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
monkeypatch.setattr(litellm, "cache", None, raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
with pytest.raises(RuntimeError, match="db gone"):
await proxy_shutdown_event()
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
async with aclosing(database_runtime.Database(prisma=fake_prisma)):
await proxy_shutdown_event()
fake_prisma.disconnect.assert_awaited_once()
assert "db gone" in caplog.text
# ---------------------------------------------------------------------------
@ -1008,7 +1021,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas
task_may_finish: Final = asyncio.Event()
task_finished: Final = asyncio.Event()
flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion
flushes: Final[
list[tuple[str, bool]]
] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion
async def slow_logging_task() -> None:
await task_may_finish.wait()
@ -1060,7 +1075,7 @@ def test_otel_global_provider_published_after_callback_init():
span was orphaned. The publish (``publish_global_otel_v2_provider``) must
therefore appear after ``_initialize_startup_logging`` in the lifespan source.
"""
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan)
source = inspect.getsource(wrapped)
init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(")
publish_pos = source.find("publish_global_otel_v2_provider(")
@ -1146,7 +1161,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
The call must follow Prisma setup so DB-backed deployments do not false-positive.
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
"""
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan)
source = inspect.getsource(wrapped)
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
warn_pos = source.find("warn_budget_without_db(")

4
tests/unit/proxy/test_proxy_server.py
View file

@ -116,7 +116,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources(
message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed"
async def run_lifespan() -> None:
async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
assert state == {"tracing_receiver": None}
assert state == {"tracing_receiver": None, "database": None}
if failure_phase == "serving":
raise RuntimeError("serving failed")
if failure_phase == "cancelled":
@ -168,7 +168,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector(
transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
) as client:
async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
assert state == {"tracing_receiver": None}
assert state == {"tracing_receiver": None, "database": None}
request: Final = asyncio.create_task(client.post(
"/admin/mcp",
headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"},

75
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -1654,6 +1654,9 @@ class MockPrisma:
async def disconnect(self):
pass
async def stop_background_tasks(self):
pass
mock_prisma = MockPrisma()
@ -1764,7 +1767,7 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database
from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError
from litellm.proxy.proxy_server import proxy_startup_event
async def unreachable():
async def unreachable(resources):
raise ConnectionError("database is down")
_, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings)
@ -1809,7 +1812,7 @@ async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values
database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with))
async def connected():
async def connected(resources):
return database
_, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY})
@ -1844,6 +1847,15 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet
class _PrismaClientWhoseDatabaseRejectsQueries:
def __init__(self) -> None:
self.disconnect_calls = 0
async def disconnect(self) -> None:
self.disconnect_calls += 1
async def stop_background_tasks(self) -> None:
pass
class _Database:
async def query_raw(self, query, *args):
raise RuntimeError("permission denied for table LiteLLM_CredentialsTable")
@ -1858,7 +1870,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc
from litellm.proxy.proxy_server import proxy_startup_event
_boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries())
prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
with (
@ -1871,6 +1884,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc
notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message]
assert len(notices) == 1
assert "Could not migrate stored values" in notices[0]
assert prisma.disconnect_calls == 1
@pytest.mark.asyncio
async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path):
from fastapi import FastAPI
from litellm.proxy.proxy_server import proxy_startup_event
_boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries())
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
with pytest.raises(RuntimeError, match="permission denied"):
async with proxy_startup_event(FastAPI()):
pass
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "")
async with proxy_startup_event(FastAPI()) as state:
assert state["database"] is None
@pytest.mark.asyncio
@ -1917,11 +1949,11 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke
class _ShutdownAwarePrisma(MockPrisma):
def __init__(self):
super().__init__()
self.stop_view_setup_task = AsyncMock()
self.stop_background_tasks = AsyncMock()
@pytest.mark.asyncio
async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path):
async def test_proxy_shutdown_stops_database_background_tasks(monkeypatch, tmp_path):
import yaml
from fastapi import FastAPI
@ -1936,11 +1968,11 @@ async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path):
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
stopped_while_serving = fake_prisma.stop_background_tasks.await_count
actual = {
"stopped_while_serving": stopped_while_serving,
"stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count,
"stopped_after_shutdown": fake_prisma.stop_background_tasks.await_count,
}
assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1}
@ -14231,6 +14263,8 @@ async def test_team_window_spend_row_carries_the_request_start_time():
def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
client = MagicMock()
client.connect = AsyncMock(side_effect=connect_error)
client.disconnect = AsyncMock()
client.stop_background_tasks = AsyncMock()
client.db.start_token_refresh_task = AsyncMock()
client.check_view_exists = AsyncMock()
client._set_spend_logs_row_count_in_proxy_state = AsyncMock()
@ -14343,6 +14377,8 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon
mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out"))
with pytest.raises(httpx.ReadTimeout):
await _run_setup_prisma_client(mock_client)
mock_client.stop_background_tasks.assert_awaited_once()
mock_client.disconnect.assert_awaited_once()
@pytest.mark.asyncio
@ -14363,6 +14399,31 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey
assert result is None
assert mock_client.start_db_health_watchdog_task.await_count == 0
assert mock_client.health_check.await_count == 0
mock_client.stop_background_tasks.assert_awaited_once()
mock_client.disconnect.assert_awaited_once()
@pytest.mark.asyncio
async def test_cancelling_prisma_setup_closes_the_client_before_publishing_it(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False")
checking: Final = asyncio.Event()
client: Final = _mock_startup_prisma_client()
async def health_check() -> None:
checking.set()
await asyncio.Event().wait()
client.health_check = health_check
setup: Final = asyncio.create_task(_run_setup_prisma_client(client))
try:
await asyncio.wait_for(checking.wait(), timeout=2)
finally:
setup.cancel()
with pytest.raises(asyncio.CancelledError):
await setup
client.stop_background_tasks.assert_awaited_once()
client.disconnect.assert_awaited_once()
async def _run_scheduled_background_jobs():

1
tests/unit/proxy/utils/prisma_and_spend/conftest.py
View file

@ -139,6 +139,7 @@ def mock_prisma_client() -> MagicMock:
client.db.is_connected = MagicMock(return_value=False)
client.db.connect = AsyncMock()
client.db.disconnect = AsyncMock()
client.db.stop_token_refresh_task = AsyncMock()
client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
client.db.execute_raw = AsyncMock()
client.db.tx = MagicMock()

156
tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py
View file

@ -14,12 +14,162 @@ from __future__ import annotations
import asyncio
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import threading
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy.utils import PrismaClient
from litellm.proxy.utils import _SHUTDOWN_RECONNECT_WAIT_SECONDS, PrismaClient
from litellm.repositories.database import Database
@pytest.mark.asyncio
@pytest.mark.parametrize("engine_dead", (False, True))
async def test_shutdown_waits_for_hung_request_reconnect_to_time_out(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, engine_dead: bool
) -> None:
monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test")
reconnecting: Final = asyncio.Event()
cancelled: Final = asyncio.Event()
async def hung_recreate(*args: object, **kwargs: object) -> None:
reconnecting.set()
try:
await asyncio.Event().wait()
finally:
cancelled.set()
prisma_client.db.recreate_prisma_client = AsyncMock(side_effect=hung_recreate)
prisma_client._engine_confirmed_dead = engine_dead
with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock:
reconnect: Final = asyncio.create_task(
prisma_client.attempt_db_reconnect(reason="request", force=True, force_recreate=True)
)
await reconnecting.wait()
closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose())
try:
await asyncio.sleep(0)
assert not closing.done()
prisma_client.db.disconnect.assert_not_awaited()
clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 1
assert await reconnect is False
await closing
finally:
reconnect.cancel()
closing.cancel()
await asyncio.gather(reconnect, closing, return_exceptions=True)
assert cancelled.is_set()
assert not prisma_client._db_reconnect_lock.locked()
prisma_client.db.disconnect.assert_awaited_once()
assert await prisma_client.attempt_db_reconnect(reason="request", force=True) is False
prisma_client.db.recreate_prisma_client.assert_awaited_once()
@pytest.mark.asyncio
async def test_shutdown_stops_waiting_for_a_held_reconnect_lock_after_the_fixed_limit(
prisma_client: PrismaClient,
) -> None:
assert _SHUTDOWN_RECONNECT_WAIT_SECONDS < prisma_client._db_watchdog_reconnect_timeout_seconds
await prisma_client._db_reconnect_lock.acquire()
with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock:
closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose())
try:
await asyncio.sleep(0)
assert not closing.done()
prisma_client.db.disconnect.assert_not_awaited()
clock.return_value = 1000.0 + _SHUTDOWN_RECONNECT_WAIT_SECONDS + 1
await closing
finally:
closing.cancel()
await asyncio.gather(closing, return_exceptions=True)
prisma_client.db.disconnect.assert_awaited_once()
prisma_client._db_reconnect_lock.release()
await asyncio.sleep(0)
assert not prisma_client._db_reconnect_lock.locked()
@pytest.mark.asyncio
async def test_shutdown_rejects_a_reconnect_already_waiting_for_the_lock(prisma_client: PrismaClient) -> None:
await prisma_client._db_reconnect_lock.acquire()
reconnect: Final = asyncio.create_task(
prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True)
)
await asyncio.sleep(0)
closing: Final = asyncio.create_task(prisma_client.stop_background_tasks())
await asyncio.sleep(0)
prisma_client._db_reconnect_lock.release()
await asyncio.wait_for(closing, timeout=2)
assert await reconnect is False
assert await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) is False
prisma_client.db.connect.assert_not_awaited()
prisma_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_shutdown_waits_for_spawned_engine_work_and_prevents_rearming(prisma_client: PrismaClient) -> None:
running: Final = asyncio.Event()
cancelled: Final = asyncio.Event()
released: Final = asyncio.Event()
finished: Final = asyncio.Event()
async def engine_work() -> None:
running.set()
try:
await asyncio.Event().wait()
finally:
cancelled.set()
await released.wait()
finished.set()
prisma_client._spawn_engine_task(engine_work())
await asyncio.wait_for(running.wait(), timeout=2)
closing: Final = asyncio.create_task(prisma_client.stop_background_tasks())
try:
await asyncio.wait_for(cancelled.wait(), timeout=2)
assert not finished.is_set()
assert not closing.done()
finally:
released.set()
await asyncio.wait_for(closing, timeout=2)
assert finished.is_set()
prisma_client._handle_writer_engine_replaced()
await asyncio.sleep(0)
assert not prisma_client._engine_tasks
@pytest.mark.asyncio
async def test_shutdown_disconnect_returns_when_an_engine_wait_thread_never_exits(
patched_prisma_import: MagicMock, mock_prisma_client: MagicMock, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False)
client: Final = PrismaClient(
database_url="postgresql://test:test@localhost:5432/test",
proxy_logging_obj=MagicMock(failure_handler=AsyncMock()),
engine_exit_wait_seconds=0.05,
)
client.db = mock_prisma_client.db
release: Final = threading.Event()
stuck: Final = threading.Thread(target=release.wait, daemon=True)
stuck.start()
client._engine_wait_threads = (stuck,)
client._closing = True
try:
await asyncio.wait_for(client.disconnect(), timeout=2)
finally:
release.set()
stuck.join(timeout=2)
client.db.disconnect.assert_awaited_once()
assert client._engine_wait_threads == ()
_PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present"

36
tests/unit/repositories/test_database.py Normal file
View file

@ -0,0 +1,36 @@
import logging
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.repositories.database import Database
@pytest.mark.asyncio
async def test_database_disconnects_and_logs_when_stopping_background_tasks_fails(
caplog: pytest.LogCaptureFixture,
) -> None:
prisma: Final = MagicMock()
prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed"))
prisma.disconnect = AsyncMock()
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
await Database(prisma=prisma).aclose()
prisma.disconnect.assert_awaited_once()
assert "watchdog failed" in caplog.text
@pytest.mark.asyncio
async def test_database_aclose_returns_and_logs_when_disconnect_fails(caplog: pytest.LogCaptureFixture) -> None:
prisma: Final = MagicMock()
prisma.stop_background_tasks = AsyncMock()
prisma.disconnect = AsyncMock(side_effect=RuntimeError("engine hung"))
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
await Database(prisma=prisma).aclose()
prisma.stop_background_tasks.assert_awaited_once()
prisma.disconnect.assert_awaited_once()
assert "engine hung" in caplog.text