From ee2cb0b6e02ceeb55d954bdae86c7c03f35bc194 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 17:21:14 -0700 Subject: [PATCH] refactor(db): keep the database handle feature-agnostic Database now holds the backend handle and its close operation only, and each feature provider builds its own repository from it, so _db never grows a field per feature. get_database checks the concrete type instead of validating a runtime Protocol on every request, and the banner audit runs as a FastAPI background task so the in-flight drain waits for it before the database closes The proxy-dependencies skill moves to the repo root .agents/skills and gains bridge rules against passing SQL text or identifiers from Python Co-Authored-By: Claude Opus 5.5 --- .../skills/proxy-dependencies/SKILL.md | 7 +++-- .../references/route-migration.md | 5 +-- .../references/runtime-ownership.md | 7 +++-- AGENTS.md | 2 +- litellm-rust/AGENTS.md | 2 -- litellm/_db.py | 27 +++------------- litellm/proxy/deps.py | 14 ++++++--- litellm/proxy/proxy_server.py | 2 +- .../user_banner_dependencies.py | 4 +-- .../user_banner_endpoints.py | 6 ++-- .../unit/proxy/proxy_server/test_lifecycle.py | 8 ++--- tests/unit/proxy/test_deps.py | 31 ++++--------------- .../test_user_banner_endpoints.py | 2 +- tests/unit/test__db.py | 2 +- 14 files changed, 44 insertions(+), 75 deletions(-) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/SKILL.md (82%) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/references/route-migration.md (85%) rename {litellm-rust/.agents => .agents}/skills/proxy-dependencies/references/runtime-ownership.md (82%) diff --git a/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md b/.agents/skills/proxy-dependencies/SKILL.md similarity index 82% rename from litellm-rust/.agents/skills/proxy-dependencies/SKILL.md rename to .agents/skills/proxy-dependencies/SKILL.md index 4d714adc48c..7ff6c54aa86 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/SKILL.md +++ b/.agents/skills/proxy-dependencies/SKILL.md @@ -5,7 +5,7 @@ description: Inject typed repositories and services into Python proxy handlers, # Proxy dependencies -- Follow the root [database-access skill](../../../../.agents/skills/database-access/SKILL.md) for queries, repository extraction and backend migration +- Follow the root [database-access skill](../database-access/SKILL.md) for queries, repository extraction and backend migration - This skill owns dependency wiring and resource lifetime, not SQL or schema ownership ## Handler contracts @@ -19,8 +19,9 @@ description: Inject typed repositories and services into Python proxy handlers, ## Database ownership -- `litellm._db.Database` exposes typed repositories and an async close operation -- `from_prisma` adopts an already initialized client without opening a second pool +- `litellm._db.Database` holds one handle per backend and an async close operation, never a field per feature +- It wraps the already initialized Prisma client without opening a second pool +- Feature modules build their repositories from that handle in their own providers, `_db` never imports a repository - FastAPI lifespan enters `_db.lifespan` and publishes its handle in request state - Drain database consumers before exiting that context, which alone closes the handle - Do not keep a parallel Prisma disconnect fallback in proxy shutdown diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md b/.agents/skills/proxy-dependencies/references/route-migration.md similarity index 85% rename from litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md rename to .agents/skills/proxy-dependencies/references/route-migration.md index d3c2aad4204..c5e2c5da878 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/references/route-migration.md +++ b/.agents/skills/proxy-dependencies/references/route-migration.md @@ -3,7 +3,7 @@ ## Existing pattern - `get_database` reads the lifespan-owned handle from request state -- A feature provider selects the narrow repository, such as `database.user_banners` +- A feature provider builds its narrow repository from the backend handle, such as `UserBannerRepository(database.prisma)` - A handler receives `UserBannerStore | None`, never a Prisma client - `UserBannerRepository` keeps the existing Prisma table, serialization and configuration-sync behavior - The banner audit provider binds the legacy audit callable at request time, without performing audit I/O @@ -17,7 +17,8 @@ - Banner GET without a database returns the default disabled banner - Banner PATCH checks the admin role before returning the existing missing-database error - Invalid stored banner data falls back to the default, but storage failures still propagate -- Banner writes finish before the success audit is scheduled +- Banner writes finish before the success audit is queued as a FastAPI background task +- Queue post-response work with `BackgroundTasks`, not bare `asyncio.create_task`, so the in-flight drain waits for it before the database closes - Do not move background work into an awaited provider merely to remove an import ## Verification diff --git a/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md b/.agents/skills/proxy-dependencies/references/runtime-ownership.md similarity index 82% rename from litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md rename to .agents/skills/proxy-dependencies/references/runtime-ownership.md index de867addc36..e2d218b5de5 100644 --- a/litellm-rust/.agents/skills/proxy-dependencies/references/runtime-ownership.md +++ b/.agents/skills/proxy-dependencies/references/runtime-ownership.md @@ -3,8 +3,8 @@ ## Current foundation - Prisma startup still owns connection setup, authentication refresh, health checks and migrations -- Lifespan enters `_db.lifespan` with the initialized client adapted by `from_prisma` -- The `Database` handle provides typed repositories and `aclose`, without exposing Prisma to handlers +- Lifespan enters `_db.lifespan` with the initialized client wrapped in `Database` +- Feature providers build repositories from the `Database` handle, so handlers never see Prisma - The existing shutdown sequence drains requests, jobs, logging, spend and audit work before disconnect - `_db.lifespan` alone closes the handle, including startup or shutdown failure - No parallel disconnect fallback belongs in proxy shutdown @@ -13,6 +13,9 @@ ## Adding a backend - Keep domain operations in repositories and conversion at the native bridge +- Add the backend as one more handle on `Database`, and let each feature provider pick which repository implementation to build +- Expose typed operations across the bridge, never a function that accepts SQL text, table names or column names from Python +- Map dynamic filters and sort keys to closed Rust enums that select static SQL, and pass every value as a bind parameter - Define who creates, borrows and closes each pool, including startup-failure cleanup - Initialize native resources in the serving process, respecting the existing native runtime's fork restrictions - Await native shutdown instead of depending on Python garbage collection diff --git a/AGENTS.md b/AGENTS.md index 69e34be00fc..24aeb7b615e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,4 +1,4 @@ -For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](litellm-rust/.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md) +For proxy dependency wiring and database lifetime changes, follow [proxy-dependencies](.agents/skills/proxy-dependencies/SKILL.md) alongside [database-access](.agents/skills/database-access/SKILL.md) Do not write comments unless they are any of: - absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index c21837858e4..e62a9463cb5 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,7 +1,5 @@ # Rust workspace rules -For proxy dependency wiring and database lifetime changes, follow [.agents/skills/proxy-dependencies/SKILL.md](.agents/skills/proxy-dependencies/SKILL.md) alongside the root database-access skill - For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm/_db.py b/litellm/_db.py index 587cbca54a5..f78a18e80f2 100644 --- a/litellm/_db.py +++ b/litellm/_db.py @@ -1,37 +1,18 @@ from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol, runtime_checkable - -from litellm.repositories.user_banner_repository import UserBannerRepository, UserBannerStore +from typing import TYPE_CHECKING if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -@runtime_checkable -class Database(Protocol): - @property - def user_banners(self) -> UserBannerStore: ... - - async def aclose(self) -> None: ... - - @dataclass(frozen=True, slots=True) -class _PrismaDatabase: - client: "PrismaClient" - - @property - def user_banners(self) -> UserBannerStore: - return UserBannerRepository(self.client) +class Database: + prisma: "PrismaClient" async def aclose(self) -> None: - await self.client.disconnect() - - -def from_prisma(client: "PrismaClient") -> Database: - """Adapt an initialized client without opening another pool.""" - return _PrismaDatabase(client) + await self.prisma.disconnect() @asynccontextmanager diff --git a/litellm/proxy/deps.py b/litellm/proxy/deps.py index 4ee0aa15c90..a10ddfedd69 100644 --- a/litellm/proxy/deps.py +++ b/litellm/proxy/deps.py @@ -1,18 +1,22 @@ from typing import Annotated, Final, TypeAlias from fastapi import Depends, Request -from pydantic import ConfigDict, TypeAdapter from litellm._db import Database -_DATABASE_ADAPTER: Final[TypeAdapter[Database | None]] = TypeAdapter( - Database | None, config=ConfigDict(arbitrary_types_allowed=True) -) _MISSING_DATABASE: Final = object() +class DatabaseStateMissingError(RuntimeError): + def __init__(self) -> None: + super().__init__("The proxy lifespan did not publish a Database in request state") + + async def get_database(request: Request) -> Database | None: - return _DATABASE_ADAPTER.validate_python(getattr(request.state, "database", _MISSING_DATABASE)) + database: Final[object] = getattr(request.state, "database", _MISSING_DATABASE) + if database is None or isinstance(database, Database): + return database + raise DatabaseStateMissingError() DatabaseDependency: TypeAlias = Annotated[Database | None, Depends(get_database)] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d2667f5b052..c2d341a5bba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1552,7 +1552,7 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener ) database: Final = await resources.enter_async_context( - _db.lifespan(None if prisma_client is None else _db.from_prisma(prisma_client)) + _db.lifespan(None if prisma_client is None else _db.Database(prisma=prisma_client)) ) await migrate_if_requested( diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py index b4016d4a226..1ad9e578e5d 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -6,7 +6,7 @@ from fastapi import Depends from litellm.models.user_banner import UserBanner from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth from litellm.proxy.deps import DatabaseDependency -from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerStore +from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository, UserBannerStore class UserBannerAudit(Protocol): @@ -14,7 +14,7 @@ class UserBannerAudit(Protocol): async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None: - return None if database is None else database.user_banners + return None if database is None else UserBannerRepository(database.prisma) async def get_user_banner_audit() -> UserBannerAudit: diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 1ba6cfed3f0..ac2c3be81e9 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py @@ -1,7 +1,6 @@ -import asyncio from typing import Annotated, Final -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException from pydantic import ConfigDict from litellm._uuid import uuid4 @@ -51,6 +50,7 @@ async def update_user_banner( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], repository: UserBannerRepositoryDependency, audit: UserBannerAuditDependency, + background_tasks: BackgroundTasks, ) -> UpdateUserBannerResponse: """ Publish, edit, or unpublish the dashboard banner. @@ -72,6 +72,6 @@ async def update_user_banner( await repository.save_banner(banner) - asyncio.create_task(audit(before, banner, user_api_key_dict)) + background_tasks.add_task(audit, before, banner, user_api_key_dict) return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index b182d3ed33f..707bd8ca237 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -118,7 +118,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 _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() fake_prisma.disconnect.assert_not_awaited() @@ -225,7 +225,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 _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -266,7 +266,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 _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -338,7 +338,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 _db.lifespan(_db.from_prisma(fake_prisma)): + async with _db.lifespan(_db.Database(prisma=fake_prisma)): await proxy_shutdown_event() diff --git a/tests/unit/proxy/test_deps.py b/tests/unit/proxy/test_deps.py index 712b206c8bb..0591064cfe2 100644 --- a/tests/unit/proxy/test_deps.py +++ b/tests/unit/proxy/test_deps.py @@ -6,15 +6,12 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from pydantic import ValidationError +from starlette.requests import Request from litellm import _db from litellm.models.user_banner import UserBanner -from litellm.proxy.deps import get_database -from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import ( - UserBannerRepositoryDependency, - get_user_banner_repository, -) +from litellm.proxy.deps import DatabaseStateMissingError, get_database +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import UserBannerRepositoryDependency def _app(database: _db.Database | None) -> FastAPI: @@ -36,7 +33,7 @@ def test_lifespan_database_is_shared_until_shutdown(): prisma: Final = MagicMock() prisma.disconnect = AsyncMock() prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) - app: Final = _app(_db.from_prisma(prisma)) + app: Final = _app(_db.Database(prisma=prisma)) with TestClient(app) as client: assert client.get("/banner").json() == UserBanner().model_dump() prisma.db.litellm_uisettings.find_unique.return_value = SimpleNamespace( @@ -63,7 +60,7 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne prisma.db.litellm_uisettings.find_unique = AsyncMock( return_value=SimpleNamespace(ui_settings={"message": "stored"}) ) - app: Final = _app(_db.from_prisma(prisma)) + app: Final = _app(_db.Database(prisma=prisma)) with TestClient(app) as client: assert client.get("/banner").json() == UserBanner(message="stored").model_dump() monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) @@ -74,22 +71,6 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne @pytest.mark.asyncio async def test_missing_lifespan_is_not_treated_as_a_disabled_database(): - from starlette.requests import Request - request: Final = Request({"type": "http", "state": {}}) - with pytest.raises(ValidationError, match="Database"): + with pytest.raises(DatabaseStateMissingError): await get_database(request) - - -@pytest.mark.asyncio -async def test_constructing_repository_does_not_touch_database(): - prisma: Final = MagicMock() - repository: Final = await get_user_banner_repository(_db.from_prisma(prisma)) - assert repository is not None - assert prisma.mock_calls == [] - - -def test_dependencies_do_not_become_openapi_parameters(): - operation: Final = _app(None).openapi()["paths"]["/banner"]["get"] - assert "parameters" not in operation - assert "requestBody" not in operation diff --git a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py index 7fdf3222869..dc3e748a5de 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_user_banner_endpoints.py @@ -58,7 +58,7 @@ def _mock_prisma(monkeypatch, record=None): mock_prisma = MagicMock() mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record) mock_prisma.db.litellm_uisettings.upsert = AsyncMock() - monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.from_prisma(mock_prisma)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=mock_prisma)) return mock_prisma diff --git a/tests/unit/test__db.py b/tests/unit/test__db.py index e6c18defd97..829bff2d014 100644 --- a/tests/unit/test__db.py +++ b/tests/unit/test__db.py @@ -12,7 +12,7 @@ from litellm import _db async def test_lifespan_closes_database_when_owner_fails(failure): prisma: Final = MagicMock() prisma.disconnect = AsyncMock() - database: Final = _db.from_prisma(prisma) + database: Final = _db.Database(prisma=prisma) async def fail(): async with _db.lifespan(database) as current: