refactor(proxy): inject deployment warning repository

This commit is contained in:
Devin AI 2026-10-27 16:47:34 +00:00 • committed by Yujong Lee
parent 3b91452385
commit dbc99e4b27
7 changed files with 52 additions and 20 deletions

View file

@ -4,6 +4,7 @@
- `get_database` reads the lifespan-owned handle from request state
- `get_repositories` builds the shared access point from that handle
- Choose each feature's replica routing once in `get_repositories`, never at call sites
- 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

View file

@ -4,6 +4,7 @@ from fastapi import Depends, Request
from litellm.repositories.container import Repositories
from litellm.repositories.database import Database
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
_MISSING_DATABASE: Final = object()
@ -25,7 +26,12 @@ DatabaseDep: TypeAlias = Annotated[Database | None, Depends(get_database)]
async def get_repositories(database: DatabaseDep) -> Repositories | None:
return None if database is None else Repositories(banners=UserBannerRepository(database.prisma))
if database is None:
return None
return Repositories(
banners=UserBannerRepository(database.prisma),
deployment_warning_ignores=DeploymentWarningIgnoreRepository(database.prisma, use_writer=True),
)
RepositoriesDep: TypeAlias = Annotated[Repositories | None, Depends(get_repositories)]

View file

@ -13,6 +13,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps import RepositoriesDep
from litellm.proxy.health_endpoints._health_endpoints import (
DISABLE_NO_REDIS_WARNING_ENV_VAR,
env_credential_login_enabled,
@ -22,7 +23,6 @@ from litellm.proxy.health_endpoints._health_endpoints import (
)
from litellm.repositories.deployment_warning_ignore_repository import (
DEPLOYMENT_WARNING_IGNORES_ROW_ID,
DeploymentWarningIgnoreRepository,
)
from litellm.types.llms.base import LiteLLMBaseModel
@ -225,12 +225,8 @@ def utc_now() -> datetime:
return datetime.now(timezone.utc)
def deployment_warning_ignore_store() -> DeploymentWarningIgnoreStore | None:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
return DeploymentWarningIgnoreRepository(prisma_client, use_writer=True)
def deployment_warning_ignore_store(repositories: RepositoriesDep) -> DeploymentWarningIgnoreStore | None:
return None if repositories is None else repositories.deployment_warning_ignores
def config_audit_logger() -> ConfigAuditLogger:

View file

@ -1,8 +1,10 @@
from dataclasses import dataclass
from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
@dataclass(frozen=True, slots=True)
class Repositories:
banners: UserBannerRepository
deployment_warning_ignores: DeploymentWarningIgnoreRepository

View file

@ -5,6 +5,7 @@ from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import MagicMock
import pytest
from fastapi import FastAPI
@ -15,6 +16,7 @@ from litellm._logging import verbose_logger
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps import get_repositories
from litellm.proxy.health_endpoints.deployment_warning_endpoints import (
WarningCondition,
config_audit_logger,
@ -28,6 +30,7 @@ from litellm.proxy.health_endpoints.deployment_warning_endpoints import (
from litellm.proxy.proxy_server import (
_dump_redacted_config, # pyright: ignore[reportPrivateUsage] # the exact redaction every config audit row goes through
)
from litellm.repositories.database import Database
from litellm.repositories.deployment_warning_ignore_repository import DEPLOYMENT_WARNING_IGNORES_ROW_ID
T0: Final = datetime(2026, 10, 14, 12, 0, tzinfo=timezone.utc)
@ -233,26 +236,24 @@ class TestCurrentWarningConditions:
class TestDeploymentWarningIgnoreStore:
def test_no_database_means_no_store(self, monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
assert deployment_warning_ignore_store() is None
@pytest.mark.asyncio
async def test_no_database_means_no_store(self):
assert deployment_warning_ignore_store(await get_repositories(None)) is None
@pytest.mark.asyncio
async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self, monkeypatch):
async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self):
def ui_settings_table(row: object) -> SimpleNamespace:
async def find_unique(where: dict[str, str]) -> object:
return row
return SimpleNamespace(litellm_uisettings=SimpleNamespace(find_unique=find_unique))
fresh = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}}
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
SimpleNamespace(
db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh))
),
fresh: Final = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}}
prisma: Final = MagicMock(
db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh))
)
store = deployment_warning_ignore_store()
repositories: Final = await get_repositories(Database(prisma=prisma))
store: Final = deployment_warning_ignore_store(repositories)
assert store is not None
assert await store.get_raw_ignores() == fresh

View file

@ -24,6 +24,10 @@ def _app(database: database_runtime.Database | None) -> FastAPI:
async def banner(repositories: RepositoriesDep) -> UserBanner:
return UserBanner() if repositories is None else await repositories.banners.get_banner()
@app.get("/ignores")
async def ignores(repositories: RepositoriesDep) -> object:
return None if repositories is None else await repositories.deployment_warning_ignores.get_raw_ignores()
return app
@ -42,6 +46,24 @@ def test_lifespan_database_is_shared_until_shutdown():
)
def test_both_repositories_use_the_lifespan_database():
prisma: Final = MagicMock()
prisma.disconnect = AsyncMock()
prisma.stop_background_tasks = AsyncMock()
prisma.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(ui_settings={"message": "stored banner"})
)
stored_ignores: Final = {"detailed_debug": {"ignored_by": "admin", "ignored_at": "2026-10-14T12:00:00Z"}}
prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(ui_settings=stored_ignores)
)
with TestClient(_app(database_runtime.Database(prisma=prisma))) as client:
assert client.get("/banner").json() == UserBanner(message="stored banner").model_dump()
assert client.get("/ignores").json() == stored_ignores
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")

View file

@ -14,6 +14,7 @@ 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.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
from tests._master_key import MASTER_KEY
@ -169,7 +170,10 @@ class TestUpdateUserBanner:
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))
repositories: Final = Repositories(
banners=UserBannerRepository(mock_prisma),
deployment_warning_ignores=DeploymentWarningIgnoreRepository(mock_prisma),
)
monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None)
monkeypatch.setitem(app.dependency_overrides, get_repositories, lambda: repositories)