From a40701ddd69711269d48b0f3ac54a220f4f1978a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 28 Oct 2026 18:37:19 -0700 Subject: [PATCH] refactor(proxy): own the lifespan database with exit stack callbacks Drop the hand-rolled lifespan(None | Database) wrapper. The proxy lifespan already holds an AsyncExitStack, so each owned client registers Database.aclose on it directly, and tests use contextlib.aclosing. Co-Authored-By: Claude Opus 5.5 --- litellm/proxy/proxy_server.py | 18 +++-- litellm/repositories/database.py | 9 --- .../unit/proxy/proxy_server/test_lifecycle.py | 13 ++-- tests/unit/repositories/test_database.py | 72 +------------------ 4 files changed, 21 insertions(+), 91 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 43640ca44b5..4a21cd91e7c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1429,11 +1429,19 @@ async def _connect_to_count_stored_values(resources: AsyncExitStack) -> Supports database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj ) if prisma_client is None: - await resources.enter_async_context(database_runtime.lifespan(database_runtime.Database(prisma=client))) + _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 @@ -1466,9 +1474,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener import json resources.callback(_forget_prisma_client) - existing_database: Final = await resources.enter_async_context( - database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client)) - ) + existing_database: Final = _own_database(resources, prisma_client) init_verbose_loggers() @@ -1578,9 +1584,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener ) ) - database: Final = existing_database or await resources.enter_async_context( - database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client)) - ) + database: Final = existing_database or _own_database(resources, prisma_client) await migrate_if_requested( environ=os.environ, diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index 99870d3e06a..c7291ad8adc 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -22,15 +22,6 @@ class Database: await self.prisma.disconnect() -@asynccontextmanager -async def lifespan(database: Database | None) -> AsyncGenerator[Database | None, None]: - try: - yield database - finally: - if database is not None: - await database.aclose() - - @asynccontextmanager async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: task: Final = asyncio.create_task(work) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 87546572e41..f25dfed196d 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -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 @@ -118,7 +119,7 @@ async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeyp monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + async with aclosing(database_runtime.Database(prisma=fake_prisma)): await proxy_shutdown_event() fake_prisma.disconnect.assert_not_awaited() @@ -226,7 +227,7 @@ 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) - async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + async with aclosing(database_runtime.Database(prisma=fake_prisma)): await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -268,7 +269,7 @@ 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())) - async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + async with aclosing(database_runtime.Database(prisma=fake_prisma)): await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -341,7 +342,7 @@ async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monk monkeypatch.setattr(litellm, "success_callback", [], raising=False) with pytest.raises(RuntimeError, match="db gone"): - async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + async with aclosing(database_runtime.Database(prisma=fake_prisma)): await proxy_shutdown_event() @@ -1017,7 +1018,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() diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index 06dc15d9f5b..008eb1124e8 100644 --- a/tests/unit/repositories/test_database.py +++ b/tests/unit/repositories/test_database.py @@ -1,76 +1,9 @@ -import asyncio from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.repositories.database import Database, lifespan - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) -async def test_lifespan_closes_database_when_owner_fails(failure): - prisma: Final = MagicMock() - prisma.stop_background_tasks = AsyncMock() - prisma.disconnect = AsyncMock() - database: Final = Database(prisma=prisma) - - async def fail(): - async with lifespan(database) as current: - assert current is database - prisma.disconnect.assert_not_awaited() - raise failure("owner failed") - - with pytest.raises(failure): - await fail() - prisma.disconnect.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_lifespan_without_database_preserves_owner_error(): - async def fail(): - async with lifespan(None) as database: - assert database is None - raise RuntimeError("owner failed") - - with pytest.raises(RuntimeError, match="owner failed"): - await fail() - - -@pytest.mark.asyncio -async def test_cancelling_owner_waits_for_database_close() -> None: - serving: Final = asyncio.Event() - closing: Final = asyncio.Event() - release_close: Final = asyncio.Event() - prisma: Final = MagicMock() - prisma.stop_background_tasks = AsyncMock() - - async def disconnect() -> None: - closing.set() - await release_close.wait() - - prisma.disconnect = AsyncMock(side_effect=disconnect) - - async def serve() -> None: - async with lifespan(Database(prisma=prisma)): - serving.set() - await asyncio.Event().wait() - - owner: Final = asyncio.create_task(serve()) - try: - await asyncio.wait_for(serving.wait(), timeout=2) - owner.cancel() - await asyncio.wait_for(closing.wait(), timeout=2) - assert not owner.done() - release_close.set() - with pytest.raises(asyncio.CancelledError): - await owner - finally: - release_close.set() - if not owner.done(): - owner.cancel() - await asyncio.gather(owner, return_exceptions=True) - prisma.disconnect.assert_awaited_once() +from litellm.repositories.database import Database @pytest.mark.asyncio @@ -80,7 +13,6 @@ async def test_database_disconnects_even_when_stopping_background_tasks_fails() prisma.disconnect = AsyncMock() with pytest.raises(RuntimeError, match="watchdog failed"): - async with lifespan(Database(prisma=prisma)): - pass + await Database(prisma=prisma).aclose() prisma.disconnect.assert_awaited_once()