From 2465cddf2297f96ac584fe4bef9f2de01159c296 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 27 Oct 2026 09:09:31 -0700 Subject: [PATCH] refactor(proxy): inject grouped database repositories --- .agents/skills/database-access/SKILL.md | 49 ++++++- .../references/python-fastapi.md | 135 ++++++++++++++++++ .../database-access/references/rust-axum.md | 43 ++++++ AGENTS.md | 2 + litellm/models/user_banner.py | 42 ++++++ litellm/proxy/deps.py | 31 ++++ .../user_banner_dependencies.py | 31 ++++ .../user_banner_endpoints.py | 96 +++---------- litellm/repositories/__init__.py | 2 + litellm/repositories/container.py | 8 ++ .../repositories/user_banner_repository.py | 20 +++ tests/unit/proxy/test_deps.py | 75 ++++++++++ .../test_user_banner_endpoints.py | 73 +++++++++- 13 files changed, 521 insertions(+), 86 deletions(-) create mode 100644 .agents/skills/database-access/references/python-fastapi.md create mode 100644 .agents/skills/database-access/references/rust-axum.md create mode 100644 litellm/models/user_banner.py create mode 100644 litellm/proxy/deps.py create mode 100644 litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py create mode 100644 litellm/repositories/container.py create mode 100644 tests/unit/proxy/test_deps.py diff --git a/.agents/skills/database-access/SKILL.md b/.agents/skills/database-access/SKILL.md index 796dfafe4f6..16dd4254a1c 100644 --- a/.agents/skills/database-access/SKILL.md +++ b/.agents/skills/database-access/SKILL.md @@ -1,6 +1,6 @@ --- name: database-access -description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope +description: Organize database repositories and proxy dependency wiring, manage connection ownership, or port Prisma operations to Rust and SQLx without expanding migration scope --- # Database access @@ -8,11 +8,44 @@ description: Add database queries, extract existing Prisma calls into repositori ## Adding new database access - Put Python database access behind an operation in `litellm/repositories` +- Group operations by domain in `Repositories`, adding fields as their callers migrate +- The container borrows existing handles and performs no database I/O during construction +- Domain protocols are optional substitution boundaries, not required for every repository - Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries -- Keep query execution and backend selection inside the repository +- Keep query execution inside the repository, with one implementation per operation and no backend switch or fallback - Prisma is fine for the implementation, adding a query does not require a Rust port - Native Rust callers use `litellm-rust/crates/db` directly +## Framework wiring + +- Python routes receive `RepositoriesDependency` and call domain operations such as `repositories.banners.get_banner()` +- Use named FastAPI providers with `Annotated[..., Depends(...)]`, ordinary helpers and jobs receive explicit arguments +- Preserve each route's no-database behavior, authorization precedence and background-task scheduling +- Read [Python and FastAPI](references/python-fastapi.md) for provider examples, lifespan state and route verification +- Read [Rust and Axum](references/rust-axum.md) when adding database access to Rust handlers + +## Connection ownership + +- Keep connection setup and shutdown separate from `Repositories`, which owns no pools +- Reuse existing handles, never connect per request or disconnect from a repository operation +- Give each pool one owner for setup and shutdown, borrowed handles never close it +- Finish requests, jobs and queued database work before closing connections +- Clean up resources opened before a startup failure, without a second disconnect fallback +- Add a backend to the existing lifecycle owner and pass its handle only to repositories that need it +- Preserve cancellation, timeout, error and transaction semantics across the native bridge + +## Changing ownership or reloading resources + +- Inventory initialization, reload, tests, background jobs and shutdown together +- Move every writer when changing authoritative storage, copying a replaceable global into state is insufficient +- A frozen container does not make its clients immutable +- Decide whether each request captures one resource generation or reads it again during processing +- Build and validate replacements before publishing them, preserving the old generation on failure +- Serialize reloads and coordinate publication with shutdown +- Drain streams and background users before retiring old resources, retaining every generation grows resource use +- Define multi-worker coordination separately from an in-process swap +- Choose mechanisms such as `ArcSwap` only after defining publication, capture and retirement + ## Asked to move existing access behind a repository - Identify the calls being moved and their transaction boundary @@ -36,17 +69,21 @@ description: Add database queries, extract existing Prisma calls into repositori - Keep typed inputs and reusable row projections in `src/` - Refresh SQLx query metadata after SQL changes - Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db` -- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository - - A query-only port can land without switching callers +- Expose typed operations through PyO3, never an API accepting SQL text, table names or column names from Python +- Map dynamic filters and sort keys to closed Rust enums selecting static SQL, and bind every value +- Switch all callers of the operation and delete its Prisma implementation in the same change +- Move each transaction as a whole, including reads, writes and locks - Verify behavior against real PostgreSQL before switching callers - Check relevant null handling, conflicts, ordering, errors, and transaction semantics -- Remove the old implementation only when no callers need it +- Repository extraction can happen in the same PR, without converting unrelated routes to dependency injection +- Read [incremental migration](references/python-fastapi.md#incremental-migration) for the caller, repository and wiring changes ## Migration context and boundaries - Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs - Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are -- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility +- Prisma and SQLx keep separate client pools to PgBouncer, which pools connections to PostgreSQL + - Size both layers together and preserve reader/writer routing and pooling-mode compatibility - PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction - Never retry a failed Rust write through Prisma when it may already have committed - Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port diff --git a/.agents/skills/database-access/references/python-fastapi.md b/.agents/skills/database-access/references/python-fastapi.md new file mode 100644 index 00000000000..26392ad9e8f --- /dev/null +++ b/.agents/skills/database-access/references/python-fastapi.md @@ -0,0 +1,135 @@ +# Python and FastAPI + +## Existing pattern + +- `get_database` reads the lifespan-owned handle from request state +- `get_repositories` builds the shared access point from that handle +- A handler receives `Repositories | None`, then uses `repositories.banners`, 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 +- SQLx can later implement the same typed operations without changing handler signatures + +## Dependency injection + +- Group operations by domain in `Repositories`, adding fields as callers migrate +- The provider builds the container from existing handles + +```python +@dataclass(frozen=True, slots=True) +class Repositories: + banners: UserBannerRepository + + +async def get_repositories(database: DatabaseDependency) -> Repositories | None: + return None if database is None else Repositories(banners=UserBannerRepository(database.prisma)) + + +RepositoriesDependency: TypeAlias = Annotated[Repositories | None, Depends(get_repositories)] +``` + +- FastAPI supplies the container to routes and caches the provider within each request +- Ordinary helpers and jobs receive explicit arguments, never unresolved `Depends` defaults + +```python +async def get_user_banner(repositories: RepositoriesDependency) -> UserBanner: + if repositories is None: + return UserBanner() + return await repositories.banners.get_banner() +``` + +## Lifespan state + +- Prisma startup owns connection setup, authentication refresh, health checks and migrations +- `Database` wraps the initialized Prisma client, without opening another pool or holding fields per feature +- FastAPI lifespan enters `litellm.repositories.database.lifespan` and publishes its handle in request state +- Starlette shallow-copies lifespan state into request state, which is distinct from `app.state` +- The database module never imports a repository, legacy globals borrow its handle without owning shutdown +- Drain requests, jobs, logging, spend and audit work before database teardown +- `litellm.repositories.database.lifespan` closes the handle on exit, including failures after entry +- When adding native resources, initialize them in the serving process, respecting runtime fork restrictions +- Await native shutdown instead of depending on Python garbage collection +- Apply the [connection ownership rules](../SKILL.md#connection-ownership) when changing the owner + +## Incremental migration + +- Port one operation while other operations stay on Prisma, without waiting for unrelated routes to adopt dependency injection +- Keep callers using `repositories.banners.get_banner()` with the same inputs and outputs +- This example ports banner reads while writes remain Prisma + - `NativeBannerAdapter` and `database.native` are proposed interfaces, not implemented by the current banner PR + - The adapter converts typed PyO3 results to `UserBanner`, preserving defaults, invalid-data handling and storage errors + - SQL stays in `crates/db/queries/`, Python never sends SQL strings + +```diff + class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): + table_name = "litellm_uisettings" + ++ def __init__(self, prisma: PrismaClient, native: NativeBannerAdapter) -> None: ++ super().__init__(prisma) ++ self._native = native ++ + async def get_banner(self) -> UserBanner: +- return parse_user_banner(await self.get_raw_settings()) ++ return await self._native.get_banner() + +- async def get_raw_settings(self) -> object: +- db_record: Final = await self.table.find_unique(where={"id": USER_BANNER_ROW_ID}) +- return db_record.ui_settings if db_record is not None else None +``` + +- Keep `save_banner` and its Prisma write unchanged +- Preserve existing inheritance only while unmigrated methods need it, new repositories use composition +- After the existing lifecycle owner initializes the native handle, update the provider + +```diff + async def get_repositories(database: DatabaseDependency) -> Repositories | None: +- return None if database is None else Repositories(banners=UserBannerRepository(database.prisma)) ++ if database is None: ++ return None ++ return Repositories( ++ banners=UserBannerRepository(prisma=database.prisma, native=database.native), ++ ) +``` + +- Both handles coexist for different operations, only the native `get_banner` implementation remains +- Switch every caller, including helpers and jobs, and delete the old query in the same PR +- Validate against real PostgreSQL and refresh SQLx metadata +- Move a complete transaction together when the query belongs to one +- Remove the Prisma handle and inherited table access after the repository's last Prisma operation moves +- Roll back the release if needed, never retry an uncertain native write through Prisma + +## Timing and errors + +- FastAPI resolves dependencies before entering the handler +- Handler exception handling cannot catch failures raised during dependency resolution +- A universal required-database dependency can override authorization and validation behavior +- 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 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 + +- Exercise real FastAPI handlers using `dependency_overrides` for external boundaries +- Keep production parsing and persistence code real +- Restore overrides after each test +- Verify complete response bodies, status codes and OpenAPI against the base +- Include combined failure conditions to pin error precedence +- Check saved values, fresh revisions, audit contents and failed-write behavior +- Check lifespan state with a running lifespan, not only a TestClient that skips startup +- Verify providers do not connect, query, disconnect or audit eagerly +- Use real PostgreSQL and live proxy requests when changing persistence behavior + +## Scope + +- Preserve direct Python call sites or migrate them explicitly +- Inventory imports, state writers, side effects and error ordering before migrating a feature +- Preserve configuration sync and serialization, and state which legacy reads remain +- Do not migrate every route in a large module just to establish the pattern +- Do not invent providers for resources without consumers + +## Framework references + +- [Dependency composition and caching](https://fastapi.tiangolo.com/tutorial/dependencies/sub-dependencies/) +- [Dependency overrides](https://fastapi.tiangolo.com/advanced/testing-dependencies/) diff --git a/.agents/skills/database-access/references/rust-axum.md b/.agents/skills/database-access/references/rust-axum.md new file mode 100644 index 00000000000..e76e231c688 --- /dev/null +++ b/.agents/skills/database-access/references/rust-axum.md @@ -0,0 +1,43 @@ +# Rust and Axum + +## Existing pattern + +- `crates/gateway` composes endpoint routers and owns server lifecycle +- Endpoint crates inject state with `.with_state(...)`, such as `crates/gateway-inference` using `Arc` +- Follow the owning crate's `AGENTS.md` before adding a database consumer + +## Database state scaffold + +- This is a proposed wiring shape, not an existing gateway database API +- The runtime supplies an initialized pool + +```rust +use axum::Router; +use sqlx::PgPool; + +#[derive(Clone)] +struct DatabaseState { + pool: PgPool, +} + +fn with_database(routes: Router, pool: PgPool) -> Router { + routes.with_state(DatabaseState { pool }) +} +``` + +- Handlers extract `State` and call typed `litellm_db` operations with the pool or transaction +- Extend existing endpoint state or use `FromRef` when other dependencies are already present +- Cloning `PgPool` shares its handle, it does not create another pool +- Keep SQL in `crates/db`, and response conversion and HTTP error mapping in the endpoint crate +- Reuse the same query APIs as Python's bridge, without creating a second SQL implementation + +## Lifetime and verification + +- Apply the [connection ownership rules](../SKILL.md#connection-ownership), finish active requests and database jobs before closing the pool +- Preserve authentication, optional-database behavior and transaction boundaries for the specific route +- Verify HTTP behavior through the owning router and query behavior against real PostgreSQL + +## Framework references + +- [Axum state and substates](https://docs.rs/axum/latest/axum/extract/struct.State.html) +- [SQLx pool ownership](https://docs.rs/sqlx/latest/sqlx/struct.Pool.html) diff --git a/AGENTS.md b/AGENTS.md index f4b1de57807..1d099c50a60 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,5 @@ +For database operations, connection ownership and proxy dependency wiring, follow [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) - used as an input for tools to read and act on. For example: diff --git a/litellm/models/user_banner.py b/litellm/models/user_banner.py new file mode 100644 index 00000000000..5470ad4dafa --- /dev/null +++ b/litellm/models/user_banner.py @@ -0,0 +1,42 @@ +from typing import Final, Literal + +from pydantic import ConfigDict, Field, model_validator + +from litellm.types.llms.base import LiteLLMBaseModel + +USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 +UserBannerSeverity = Literal["info", "warning", "error"] + + +class UserBannerUpdate(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + enabled: bool = Field( + default=False, + description="If true, the banner is shown to all authenticated dashboard users.", + ) + message: str = Field( + default="", + max_length=USER_BANNER_MAX_MESSAGE_LENGTH, + description="Banner text shown to dashboard users. Markdown is supported.", + ) + severity: UserBannerSeverity = Field( + default="info", + description="Visual style of the banner.", + ) + + @model_validator(mode="after") + def _require_message_when_enabled(self) -> "UserBannerUpdate": + if self.enabled and not self.message.strip(): + raise ValueError("message must be non-empty when the banner is enabled") + return self + + +class UserBanner(UserBannerUpdate): + revision: str = Field( + default="", + description=( + "Server-stamped opaque publish identity; a fresh value is generated on every " + "update so clients re-surface dismissed banners on republish." + ), + ) diff --git a/litellm/proxy/deps.py b/litellm/proxy/deps.py new file mode 100644 index 00000000000..ccb5c990732 --- /dev/null +++ b/litellm/proxy/deps.py @@ -0,0 +1,31 @@ +from typing import Annotated, Final, TypeAlias + +from fastapi import Depends, Request + +from litellm.repositories.container import Repositories +from litellm.repositories.database import Database +from litellm.repositories.user_banner_repository import UserBannerRepository + +_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: + 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)] + + +async def get_repositories(database: DatabaseDependency) -> Repositories | None: + return None if database is None else Repositories(banners=UserBannerRepository(database.prisma)) + + +RepositoriesDependency: TypeAlias = Annotated[Repositories | None, Depends(get_repositories)] diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py new file mode 100644 index 00000000000..0b8f2afe80c --- /dev/null +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -0,0 +1,31 @@ +from collections.abc import Coroutine +from typing import Annotated, Protocol, TypeAlias + +from fastapi import Depends + +from litellm.models.user_banner import UserBanner +from litellm.proxy._types import LitellmTableNames, UserAPIKeyAuth +from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID + + +class UserBannerAudit(Protocol): + def __call__(self, before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> Coroutine[None, None, None]: ... + + +async def get_user_banner_audit() -> UserBannerAudit: + from litellm.proxy.proxy_server import create_config_audit_log + + async def audit(before: UserBanner, after: UserBanner, user: UserAPIKeyAuth) -> None: + await create_config_audit_log( + param_name=USER_BANNER_ROW_ID, + action="updated", + before_value=before.model_dump(), + after_value=after.model_dump(), + user_api_key_dict=user, + table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, + ) + + return audit + + +UserBannerAuditDependency: TypeAlias = Annotated[UserBannerAudit, Depends(get_user_banner_audit)] diff --git a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py index 56339202d46..42d930760dd 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_endpoints.py @@ -1,88 +1,43 @@ -import asyncio -import json -from typing import Annotated, Final, Literal +from typing import Annotated, Final -from fastapi import APIRouter, Depends, HTTPException -from pydantic import Field, ValidationError, model_validator +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException +from pydantic import ConfigDict from litellm._uuid import uuid4 -from litellm.proxy._types import LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth +from litellm.models.user_banner import UserBanner, UserBannerUpdate +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.repositories.user_banner_repository import USER_BANNER_ROW_ID, UserBannerRepository +from litellm.proxy.deps import RepositoriesDependency +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import ( + UserBannerAuditDependency, +) from litellm.types.llms.base import LiteLLMBaseModel router: Final = APIRouter() -USER_BANNER_MAX_MESSAGE_LENGTH: Final = 4000 - -UserBannerSeverity = Literal["info", "warning", "error"] - - -class UserBannerUpdate(LiteLLMBaseModel): - enabled: bool = Field( - default=False, - description="If true, the banner is shown to all authenticated dashboard users.", - ) - message: str = Field( - default="", - max_length=USER_BANNER_MAX_MESSAGE_LENGTH, - description="Banner text shown to dashboard users. Markdown is supported.", - ) - severity: UserBannerSeverity = Field( - default="info", - description="Visual style of the banner.", - ) - - @model_validator(mode="after") - def _require_message_when_enabled(self) -> "UserBannerUpdate": - if self.enabled and not self.message.strip(): - raise ValueError("message must be non-empty when the banner is enabled") - return self - - -class UserBanner(UserBannerUpdate): - revision: str = Field( - default="", - description=( - "Server-stamped opaque publish identity; a fresh value is generated on every " - "update so clients re-surface dismissed banners on republish." - ), - ) - class UpdateUserBannerResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + message: str banner: UserBanner -def parse_user_banner(raw_settings: object) -> UserBanner: - if raw_settings is None: - return UserBanner() - try: - parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings - return UserBanner.model_validate(parsed) - except (json.JSONDecodeError, ValidationError): - return UserBanner() - - @router.get( "/get/user_banner", tags=["UI Settings"], dependencies=[Depends(user_api_key_auth)], response_model=UserBanner, ) -async def get_user_banner() -> UserBanner: +async def get_user_banner(repositories: RepositoriesDependency) -> UserBanner: """ Get the admin-published dashboard banner. Readable by any authenticated user; rendered on every dashboard page. """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: + if repositories is None: return UserBanner() - raw_settings: Final = await UserBannerRepository(prisma_client).get_raw_settings() - return parse_user_banner(raw_settings) + return await repositories.banners.get_banner() @router.patch( @@ -93,21 +48,21 @@ async def get_user_banner() -> UserBanner: async def update_user_banner( banner_update: UserBannerUpdate, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + repositories: RepositoriesDependency, + audit: UserBannerAuditDependency, + background_tasks: BackgroundTasks, ) -> UpdateUserBannerResponse: """ Publish, edit, or unpublish the dashboard banner. Only proxy admins are allowed to modify it. """ - from litellm.proxy.proxy_server import create_config_audit_log, prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="Only proxy admins can update the user banner.") - if prisma_client is None: + if repositories is None: raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.") - repository: Final = UserBannerRepository(prisma_client) - before: Final = parse_user_banner(await repository.get_raw_settings()) + before: Final = await repositories.banners.get_banner() banner: Final = UserBanner( enabled=banner_update.enabled, message=banner_update.message, @@ -115,17 +70,8 @@ async def update_user_banner( revision=uuid4().hex, ) - await repository.upsert_settings(json.dumps(banner.model_dump())) + await repositories.banners.save_banner(banner) - asyncio.create_task( - create_config_audit_log( - param_name=USER_BANNER_ROW_ID, - action="updated", - before_value=before.model_dump(), - after_value=banner.model_dump(), - user_api_key_dict=user_api_key_dict, - table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME, - ) - ) + background_tasks.add_task(audit, before, banner, user_api_key_dict) return UpdateUserBannerResponse(message="User banner updated successfully", banner=banner) diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index f090911c416..bf99ae06de2 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -5,6 +5,7 @@ Repository classes for database operations. from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.container import Repositories from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.model_repository import ModelRepository @@ -141,6 +142,7 @@ __all__ = [ "ProjectRepository", "PromptRepository", "ReadOnlyTable", + "Repositories", "SSOConfigRepository", "SearchToolsRepository", "SkillsRepository", diff --git a/litellm/repositories/container.py b/litellm/repositories/container.py new file mode 100644 index 00000000000..a25f065ec34 --- /dev/null +++ b/litellm/repositories/container.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + +from litellm.repositories.user_banner_repository import UserBannerRepository + + +@dataclass(frozen=True, slots=True) +class Repositories: + banners: UserBannerRepository diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index 4e113dff988..bc25d3af1fe 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/litellm/repositories/user_banner_repository.py @@ -1,5 +1,9 @@ +import json from typing import TYPE_CHECKING, Final +from pydantic import ValidationError + +from litellm.models.user_banner import UserBanner from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: @@ -8,9 +12,25 @@ if TYPE_CHECKING: USER_BANNER_ROW_ID: Final = "user_banner" +def parse_user_banner(raw_settings: object) -> UserBanner: + if raw_settings is None: + return UserBanner() + try: + parsed: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings + return UserBanner.model_validate(parsed) + except (json.JSONDecodeError, ValidationError): + return UserBanner() + + class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" + async def get_banner(self) -> UserBanner: + return parse_user_banner(await self.get_raw_settings()) + + async def save_banner(self, banner: UserBanner) -> None: + await self.upsert_settings(json.dumps(banner.model_dump())) + async def get_raw_settings(self) -> object: db_record: Final = await self.table.find_unique(where={"id": USER_BANNER_ROW_ID}) return db_record.ui_settings if db_record is not None else None diff --git a/tests/unit/proxy/test_deps.py b/tests/unit/proxy/test_deps.py new file mode 100644 index 00000000000..e7b45c9c296 --- /dev/null +++ b/tests/unit/proxy/test_deps.py @@ -0,0 +1,75 @@ +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.requests import Request + +from litellm.models.user_banner import UserBanner +from litellm.proxy.deps import DatabaseStateMissingError, RepositoriesDependency, get_database +from litellm.repositories import database as database_runtime + + +def _app(database: database_runtime.Database | None) -> FastAPI: + @asynccontextmanager + async def lifespan(app: FastAPI): + async with database_runtime.lifespan(database) as owned_database: + yield {"database": owned_database} + + app: Final = FastAPI(lifespan=lifespan) + + @app.get("/banner") + async def banner(repositories: RepositoriesDependency) -> UserBanner: + return UserBanner() if repositories is None else await repositories.banners.get_banner() + + return app + + +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(database_runtime.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( + ui_settings={"enabled": True, "message": "published", "revision": "revision"} + ) + assert ( + client.get("/banner").json() + == UserBanner(enabled=True, message="published", revision="revision").model_dump() + ) + prisma.disconnect.assert_not_awaited() + prisma.disconnect.assert_awaited_once() + + +def test_no_database_state_preserves_optional_repository(): + with TestClient(_app(None)) as client: + response: Final = client.get("/banner") + assert response.status_code == 200 + assert response.json() == UserBanner().model_dump() + + +def test_dependency_override_changes_next_request_without_changing_lifespan_owner(monkeypatch): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace(ui_settings={"message": "stored"}) + ) + app: Final = _app(database_runtime.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) + assert client.get("/banner").json() == UserBanner().model_dump() + prisma.disconnect.assert_not_awaited() + prisma.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_missing_lifespan_is_not_treated_as_a_disabled_database(): + request: Final = Request({"type": "http", "state": {}}) + with pytest.raises(DatabaseStateMissingError): + await get_database(request) 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 c014ffd04c6..306b1b3eede 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 @@ -1,5 +1,7 @@ +import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -7,7 +9,12 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps import get_database, get_repositories from litellm.proxy.proxy_server import app +from litellm.proxy.ui_crud_endpoints.user_banner_dependencies import get_user_banner_audit +from litellm.repositories.container import Repositories +from litellm.repositories.database import Database +from litellm.repositories.user_banner_repository import UserBannerRepository from tests._master_key import MASTER_KEY client = TestClient(app) @@ -46,7 +53,7 @@ def internal_user_auth(): @pytest.fixture def mock_audit_log(monkeypatch): audit_mock = AsyncMock(return_value=None) - monkeypatch.setattr("litellm.proxy.proxy_server.create_config_audit_log", audit_mock) + monkeypatch.setitem(app.dependency_overrides, get_user_banner_audit, lambda: audit_mock) return audit_mock @@ -54,7 +61,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.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: Database(prisma=mock_prisma)) return mock_prisma @@ -94,21 +101,77 @@ class TestGetUserBanner: assert response.json() == DISABLED_BANNER def test_no_database_returns_disabled_banner(self, admin_auth, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) response = client.get("/get/user_banner") assert response.status_code == 200 assert response.json() == DISABLED_BANNER class TestUpdateUserBanner: + def test_storage_finishes_before_audit_runs(self, admin_auth, monkeypatch, mock_audit_log): + events = [] + prisma = _mock_prisma(monkeypatch) + + async def save(**kwargs): + await asyncio.sleep(0) + events.append("saved") + + async def audit(before, after, user): + events.append("audited") + + prisma.db.litellm_uisettings.upsert.side_effect = save + mock_audit_log.side_effect = audit + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + assert events == ["saved", "audited"] + + @pytest.mark.parametrize( + ("role", "status", "detail"), + [ + (LitellmUserRoles.INTERNAL_USER, 403, "Only proxy admins can update the user banner."), + (LitellmUserRoles.PROXY_ADMIN, 500, "Database not connected. Please connect a database."), + ], + ) + def test_no_database_preserves_authorization_precedence(self, monkeypatch, role, status, detail): + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override(role)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == status + assert response.json() == {"detail": detail} + + def test_failed_write_does_not_schedule_success_audit(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch) + prisma.db.litellm_uisettings.upsert.side_effect = RuntimeError("storage unavailable") + with pytest.raises(RuntimeError, match="storage unavailable"): + client.patch("/update/user_banner", json=PUBLISH_BODY) + mock_audit_log.assert_not_called() + + def test_audit_receives_previous_and_saved_banner(self, admin_auth, monkeypatch, mock_audit_log): + prisma = _mock_prisma(monkeypatch, record=SimpleNamespace(ui_settings=json.dumps(PUBLISHED_BANNER))) + response = client.patch("/update/user_banner", json=PUBLISH_BODY) + assert response.status_code == 200 + before, after, user = mock_audit_log.call_args.args + assert before.model_dump() == PUBLISHED_BANNER + assert after.model_dump() == response.json()["banner"] + assert user.user_id == "test-user" + assert ( + json.loads(prisma.db.litellm_uisettings.upsert.await_args.kwargs["data"]["create"]["ui_settings"]) + == after.model_dump() + ) + def test_rejects_non_admin(self, internal_user_auth, monkeypatch): mock_prisma = _mock_prisma(monkeypatch) response = client.patch("/update/user_banner", json=PUBLISH_BODY) assert response.status_code == 403 mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited() - def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log): - mock_prisma = _mock_prisma(monkeypatch, record=None) + @pytest.mark.parametrize("override_repositories", [False, True]) + def test_persists_and_round_trips(self, admin_auth, monkeypatch, mock_audit_log, override_repositories: bool): + mock_prisma: Final = _mock_prisma(monkeypatch, record=None) + if override_repositories: + repositories: Final = Repositories(banners=UserBannerRepository(mock_prisma)) + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) + monkeypatch.setitem(app.dependency_overrides, get_repositories, lambda: repositories) response = client.patch("/update/user_banner", json=PUBLISH_BODY) assert response.status_code == 200