From cebc7e6dd7ff52e936e1c53cc4021ab72bf1f7aa Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 17:39:06 -0700 Subject: [PATCH] wip --- litellm-rust/Cargo.lock | 1 + ...ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json | 26 ++++ ...8dc13486a1f56d03400f54edf3f25be95cccc.json | 14 +++ .../queries/user_banner/find_user_banner.sql | 3 + .../queries/user_banner/save_user_banner.sql | 5 + litellm-rust/crates/db/src/lib.rs | 2 + litellm-rust/crates/db/src/user_banner.rs | 114 ++++++++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 13 +- .../crates/python-bridge/src/database.rs | 68 +++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 4 + litellm/_db.py | 8 +- litellm/constants.py | 1 + litellm/proxy/proxy_server.py | 17 ++- .../user_banner_dependencies.py | 13 +- .../repositories/user_banner_repository.py | 14 +++ litellm/rust_bridge/_native.pyi | 9 ++ litellm/rust_bridge/database.py | 33 +++++ .../test_user_banner_endpoints.py | 41 +++++++ tests/unit/test__db.py | 14 +++ 19 files changed, 391 insertions(+), 9 deletions(-) create mode 100644 litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json create mode 100644 litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json create mode 100644 litellm-rust/crates/db/queries/user_banner/find_user_banner.sql create mode 100644 litellm-rust/crates/db/queries/user_banner/save_user_banner.sql create mode 100644 litellm-rust/crates/db/src/user_banner.rs create mode 100644 litellm-rust/crates/python-bridge/src/database.rs create mode 100644 litellm/rust_bridge/database.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 32ae431f863..33a7c7aa30d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4229,6 +4229,7 @@ dependencies = [ "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core-utils", + "litellm-db", "litellm-host", "litellm-host-python", "litellm-http", diff --git a/litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json b/litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json new file mode 100644 index 00000000000..8b49a9c5e82 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae.json @@ -0,0 +1,26 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT ui_settings\nFROM \"LiteLLM_UISettings\"\nWHERE id = 'user_banner'\n", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "ui_settings", + "type_info": "Jsonb", + "origin": { + "Table": { + "table": "\"LiteLLM_UISettings\"", + "name": "ui_settings" + } + } + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + false + ] + }, + "hash": "136362927bf4b07894c7098e4f6ebd0f1c8d415ecdb3eb9f1a698e879dd270ae" +} diff --git a/litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json b/litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json new file mode 100644 index 00000000000..88f62dd8e06 --- /dev/null +++ b/litellm-rust/crates/db/.sqlx/query-c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO \"LiteLLM_UISettings\" (id, ui_settings, updated_at)\nVALUES ('user_banner', $1, CURRENT_TIMESTAMP)\nON CONFLICT (id) DO UPDATE\nSET ui_settings = EXCLUDED.ui_settings,\n updated_at = EXCLUDED.updated_at\n", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Jsonb" + ] + }, + "nullable": [] + }, + "hash": "c8da5047ab370f06f239ef8d5798dc13486a1f56d03400f54edf3f25be95cccc" +} diff --git a/litellm-rust/crates/db/queries/user_banner/find_user_banner.sql b/litellm-rust/crates/db/queries/user_banner/find_user_banner.sql new file mode 100644 index 00000000000..6833855a92a --- /dev/null +++ b/litellm-rust/crates/db/queries/user_banner/find_user_banner.sql @@ -0,0 +1,3 @@ +SELECT ui_settings +FROM "LiteLLM_UISettings" +WHERE id = 'user_banner' diff --git a/litellm-rust/crates/db/queries/user_banner/save_user_banner.sql b/litellm-rust/crates/db/queries/user_banner/save_user_banner.sql new file mode 100644 index 00000000000..193c08d409c --- /dev/null +++ b/litellm-rust/crates/db/queries/user_banner/save_user_banner.sql @@ -0,0 +1,5 @@ +INSERT INTO "LiteLLM_UISettings" (id, ui_settings, updated_at) +VALUES ('user_banner', $1, CURRENT_TIMESTAMP) +ON CONFLICT (id) DO UPDATE +SET ui_settings = EXCLUDED.ui_settings, + updated_at = EXCLUDED.updated_at diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index 541f97c7db9..501b890e568 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -13,6 +13,7 @@ mod failure; mod spend; #[cfg(test)] mod test_support; +mod user_banner; pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow}; pub use end_user::{EndUser, RateLimits, find_end_user}; @@ -24,3 +25,4 @@ pub use spend::{ begin, insert_spend_logs, lock_team, upsert_daily_end_user_spend, upsert_daily_spend, upsert_daily_team_spend, upsert_daily_user_spend, }; +pub use user_banner::{find_user_banner, save_user_banner}; diff --git a/litellm-rust/crates/db/src/user_banner.rs b/litellm-rust/crates/db/src/user_banner.rs new file mode 100644 index 00000000000..a1f0207fdfd --- /dev/null +++ b/litellm-rust/crates/db/src/user_banner.rs @@ -0,0 +1,114 @@ +use serde_json::Value; +use sqlx::PgExecutor; + +use crate::Error; + +pub async fn find_user_banner(executor: impl PgExecutor<'_>) -> Result, Error> { + Ok( + sqlx::query_file_scalar!("queries/user_banner/find_user_banner.sql") + .fetch_optional(executor) + .await?, + ) +} + +pub async fn save_user_banner( + executor: impl PgExecutor<'_>, + settings: &Value, +) -> Result<(), Error> { + sqlx::query_file!("queries/user_banner/save_user_banner.sql", settings) + .execute(executor) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + + use rstest::rstest; + use serde_json::{Value, json}; + use sqlx::PgPool; + use sqlx::types::chrono::NaiveDateTime; + + use super::{find_user_banner, save_user_banner}; + use crate::test_support::database_pool; + + async fn clear(tx: &mut sqlx::PgConnection) -> Result<(), sqlx::Error> { + #[expect( + clippy::disallowed_methods, + reason = "Fixture cleanup is outside the compile-checked query set" + )] + sqlx::query(r#"DELETE FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#) + .execute(tx) + .await?; + Ok(()) + } + + async fn stored_row( + tx: &mut sqlx::PgConnection, + ) -> Result<(Value, NaiveDateTime, NaiveDateTime), sqlx::Error> { + #[expect( + clippy::disallowed_methods, + reason = "Fixture reads are outside the compile-checked query set" + )] + sqlx::query_as( + r#"SELECT ui_settings, created_at, updated_at FROM "LiteLLM_UISettings" WHERE id = 'user_banner'"#, + ) + .fetch_one(tx) + .await + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn missing_banner_reads_as_none( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let mut tx = database_pool?.begin().await?; + clear(&mut tx).await?; + + assert_eq!(find_user_banner(&mut *tx).await?, None); + Ok(()) + } + + #[rstest] + #[case::first_publish(None)] + #[case::republish(Some(json!({"enabled": false, "message": "old", "severity": "info", "revision": "a"})))] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn save_replaces_the_stored_banner_and_stamps_updated_at( + #[case] previous: Option, + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let mut tx = database_pool?.begin().await?; + clear(&mut tx).await?; + if let Some(previous) = &previous { + save_user_banner(&mut *tx, previous).await?; + } + let banner = json!({"enabled": true, "message": "maintenance at 5pm", "severity": "warning", "revision": "b"}); + + save_user_banner(&mut *tx, &banner).await?; + + assert_eq!(find_user_banner(&mut *tx).await?, Some(banner.clone())); + let (stored, created_at, updated_at) = stored_row(&mut tx).await?; + assert_eq!(stored, banner); + assert!(updated_at >= created_at); + Ok(()) + } + + #[rstest] + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn banner_text_is_stored_as_data_not_sql( + #[future(awt)] database_pool: Result>, + ) -> Result<(), Box> { + let mut tx = database_pool?.begin().await?; + clear(&mut tx).await?; + let banner = json!({"message": r#"'); DROP TABLE "LiteLLM_UISettings"; --"#}); + + save_user_banner(&mut *tx, &banner).await?; + + assert_eq!(find_user_banner(&mut *tx).await?, Some(banner)); + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6fc05c43146..0a7a072d171 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,18 +19,14 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] -fancy-regex.workspace = true litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true -bytes.workspace = true -futures-util.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true -serde.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true litellm-callbacks-legacy-python.workspace = true @@ -41,6 +37,7 @@ litellm-inference-messages.workspace = true litellm-inference-chat.workspace = true litellm-inference-ocr.workspace = true litellm-core-utils.workspace = true +litellm-db.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "hashicorp", "cyberark"] } @@ -49,6 +46,11 @@ litellm-llms-types.workspace = true litellm-router-types.workspace = true litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } + +serde.workspace = true +futures-util.workspace = true +bytes.workspace = true +fancy-regex.workspace = true pyo3.workspace = true prost.workspace = true pyo3-async-runtimes.workspace = true @@ -61,9 +63,10 @@ thiserror.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] -tracing.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-secrets-aws.workspace = true + +tracing.workspace = true serde.workspace = true serde_with.workspace = true criterion.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/database.rs b/litellm-rust/crates/python-bridge/src/database.rs new file mode 100644 index 00000000000..251e4270ebe --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/database.rs @@ -0,0 +1,68 @@ +use std::error::Error as StdError; + +use litellm_db::Error; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; +use serde_json::Value; +use sqlx::{PgPool, postgres::PgPoolOptions}; + +fn map_error(error: Error) -> PyErr { + let chain: Vec = + std::iter::successors(Some(&error as &dyn StdError), |error| error.source()) + .map(ToString::to_string) + .collect(); + PyRuntimeError::new_err(chain.join(": ")) +} + +#[pyclass(frozen)] +pub struct NativeDatabase { + pool: PgPool, +} + +#[pymethods] +impl NativeDatabase { + #[staticmethod] + fn connect(py: Python<'_>, url: String, max_connections: u32) -> PyResult> { + crate::execution::run_async_value(py, async move { + let pool = PgPoolOptions::new() + .max_connections(max_connections) + .connect(&url) + .await + .map_err(|error| map_error(error.into()))?; + Ok(Self { pool }) + }) + } + + fn find_user_banner<'py>(&self, py: Python<'py>) -> PyResult> { + let pool = self.pool.clone(); + crate::execution::run_async( + py, + async move { litellm_db::find_user_banner(&pool).await }, + map_error, + ) + } + + fn save_user_banner<'py>( + &self, + py: Python<'py>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] settings: Value, + ) -> PyResult> { + let pool = self.pool.clone(); + crate::execution::run_async( + py, + async move { litellm_db::save_user_banner(&pool, &settings).await }, + map_error, + ) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let pool = self.pool.clone(); + crate::execution::run_async( + py, + async move { + pool.close().await; + Ok::<(), Error>(()) + }, + map_error, + ) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 12827b76ff3..fc7b5545545 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -2,6 +2,7 @@ mod cache; mod callable; mod coercion; mod credentials; +mod database; mod diagnostics; mod errors; mod execution; @@ -20,6 +21,8 @@ mod _native { use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; use pyo3::{prelude::*, types::PyModule}; + #[pymodule_export] + use crate::database::NativeDatabase; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -87,6 +90,7 @@ mod tests { Python::attach(|py| { let mut expected = vec![ "NativeCacheHandle", + "NativeDatabase", "RustBridgeDeclined", "RustUpstreamError", "ForkedAfterNativeRuntimeStarted", diff --git a/litellm/_db.py b/litellm/_db.py index f78a18e80f2..c064cbf2fdb 100644 --- a/litellm/_db.py +++ b/litellm/_db.py @@ -5,14 +5,20 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient + from litellm.rust_bridge.database import NativeDatabase @dataclass(frozen=True, slots=True) class Database: prisma: "PrismaClient" + native: "NativeDatabase | None" = None async def aclose(self) -> None: - await self.prisma.disconnect() + try: + if self.native is not None: + await self.native.close() + finally: + await self.prisma.disconnect() @asynccontextmanager diff --git a/litellm/constants.py b/litellm/constants.py index d1daf2ee2d1..dbd9351a4e8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2350,3 +2350,4 @@ HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset( ) DEFAULT_TOOL_LOOP_MAX_ROUNDS: Final = 20 +NATIVE_DATABASE_MAX_CONNECTIONS: Final = 5 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2d341a5bba..9b9b28709e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -86,6 +86,7 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + NATIVE_DATABASE_MAX_CONNECTIONS, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) @@ -214,6 +215,7 @@ if TYPE_CHECKING: from litellm.integrations.opentelemetry import OpenTelemetry from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager + from litellm.rust_bridge.database import NativeDatabase Span = _Span | Any else: @@ -1420,6 +1422,15 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: return client.writer_db +async def _connect_native_database() -> "NativeDatabase | None": + database_url: Final = get_secret_str("DATABASE_URL") + if database_url is None or get_secret_bool("LITELLM_NATIVE_DATABASE") is not True: + return None + from litellm.rust_bridge.database import connect_native_database + + return await connect_native_database(database_url, NATIVE_DATABASE_MAX_CONNECTIONS) + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: @@ -1552,7 +1563,11 @@ 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.Database(prisma=prisma_client)) + _db.lifespan( + None + if prisma_client is None + else _db.Database(prisma=prisma_client, native=await _connect_native_database()) + ) ) 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 1ad9e578e5d..895ee54cba9 100644 --- a/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py +++ b/litellm/proxy/ui_crud_endpoints/user_banner_dependencies.py @@ -6,7 +6,12 @@ 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, UserBannerRepository, UserBannerStore +from litellm.repositories.user_banner_repository import ( + USER_BANNER_ROW_ID, + NativeUserBannerRepository, + UserBannerRepository, + UserBannerStore, +) class UserBannerAudit(Protocol): @@ -14,7 +19,11 @@ class UserBannerAudit(Protocol): async def get_user_banner_repository(database: DatabaseDependency) -> UserBannerStore | None: - return None if database is None else UserBannerRepository(database.prisma) + if database is None: + return None + if database.native is not None: + return NativeUserBannerRepository(database.native) + return UserBannerRepository(database.prisma) async def get_user_banner_audit() -> UserBannerAudit: diff --git a/litellm/repositories/user_banner_repository.py b/litellm/repositories/user_banner_repository.py index e8298f4aa0d..ffdaecdf14c 100644 --- a/litellm/repositories/user_banner_repository.py +++ b/litellm/repositories/user_banner_repository.py @@ -1,4 +1,5 @@ import json +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol from pydantic import ValidationError @@ -9,6 +10,8 @@ from litellm.repositories.table_repositories import PrismaTableRepository if TYPE_CHECKING: from prisma import models as prisma_models # noqa: F401 # resolved only from the quoted base-class subscript below + from litellm.rust_bridge.database import NativeDatabase + USER_BANNER_ROW_ID: Final = "user_banner" @@ -28,6 +31,17 @@ def parse_user_banner(raw_settings: object) -> UserBanner: return UserBanner() +@dataclass(frozen=True, slots=True) +class NativeUserBannerRepository: + native: "NativeDatabase" + + async def get_banner(self) -> UserBanner: + return parse_user_banner(await self.native.find_user_banner()) + + async def save_banner(self, banner: UserBanner) -> None: + await self.native.save_user_banner(banner.model_dump()) + + class UserBannerRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]): table_name = "litellm_uisettings" diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7364a1d72f2..f073620a310 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -25,6 +25,14 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +@final +class NativeDatabase: + @staticmethod + def connect(url: str, max_connections: int) -> Future[NativeDatabase]: ... + def close(self) -> Future[None]: ... + def find_user_banner(self) -> Future[JsonValue]: ... + def save_user_banner(self, settings: Mapping[str, JsonValue]) -> Future[None]: ... + @final class NativeDiagnosticProcessor: def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ... @@ -207,6 +215,7 @@ __all__ = [ "HuggingFaceEncoding", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", + "NativeDatabase", "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", diff --git a/litellm/rust_bridge/database.py b/litellm/rust_bridge/database.py new file mode 100644 index 00000000000..91c1f7d64bc --- /dev/null +++ b/litellm/rust_bridge/database.py @@ -0,0 +1,33 @@ +from collections.abc import Awaitable, Mapping +from typing import Final, Protocol, runtime_checkable + +from pydantic import ConfigDict, TypeAdapter + +from litellm.rust_bridge.loader import get_native_bridge + + +class NativeDatabase(Protocol): + def find_user_banner(self) -> Awaitable[object]: ... + def save_user_banner(self, settings: Mapping[str, object]) -> Awaitable[None]: ... + def close(self) -> Awaitable[None]: ... + + +class NativeDatabaseFactory(Protocol): + def connect(self, url: str, max_connections: int) -> Awaitable[NativeDatabase]: ... + + +@runtime_checkable +class NativeDatabaseBridge(Protocol): + NativeDatabase: NativeDatabaseFactory + + +_NATIVE: Final[TypeAdapter[NativeDatabaseBridge]] = TypeAdapter( + NativeDatabaseBridge, config=ConfigDict(arbitrary_types_allowed=True) +) + + +async def connect_native_database(url: str, max_connections: int) -> NativeDatabase: + bridge: Final = get_native_bridge() + if bridge is None: + raise RuntimeError("The native database requires the Rust extension") + return await _NATIVE.validate_python(bridge).NativeDatabase.connect(url, max_connections) 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 dc3e748a5de..a7e01252539 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 @@ -229,3 +229,44 @@ class TestUpdateUserBanner: response = client.patch("/update/user_banner", json=payload) assert response.status_code == 422 mock_prisma.db.litellm_uisettings.upsert.assert_not_awaited() + + +class _NativeDatabase: + def __init__(self, stored=None): + self.stored = stored + + async def find_user_banner(self): + return self.stored + + async def save_user_banner(self, settings): + self.stored = settings + + +class TestNativeUserBannerBackend: + @pytest.fixture + def backends(self, monkeypatch): + prisma = MagicMock() + banners = _NativeDatabase() + monkeypatch.setitem(app.dependency_overrides, get_database, lambda: _db.Database(prisma=prisma, native=banners)) + return prisma, banners + + def test_published_banner_round_trips_through_the_native_backend(self, admin_auth, mock_audit_log, backends): + prisma, banners = backends + published = client.patch("/update/user_banner", json=PUBLISH_BODY).json()["banner"] + + assert banners.stored == published + assert client.get("/get/user_banner").json() == published + assert prisma.mock_calls == [] + + @pytest.mark.parametrize( + "stored", [json.dumps(PUBLISHED_BANNER), PUBLISHED_BANNER], ids=["prisma-string", "object"] + ) + def test_native_backend_reads_rows_written_by_either_backend(self, admin_auth, backends, stored): + _, banners = backends + banners.stored = stored + assert client.get("/get/user_banner").json() == PUBLISHED_BANNER + + def test_native_backend_treats_invalid_stored_banner_as_disabled(self, admin_auth, backends): + _, banners = backends + banners.stored = {"enabled": True, "message": ""} + assert client.get("/get/user_banner").json() == DISABLED_BANNER diff --git a/tests/unit/test__db.py b/tests/unit/test__db.py index 829bff2d014..35b1ec9623a 100644 --- a/tests/unit/test__db.py +++ b/tests/unit/test__db.py @@ -34,3 +34,17 @@ async def test_lifespan_without_database_preserves_owner_error(): with pytest.raises(RuntimeError, match="owner failed"): await fail() + + +@pytest.mark.asyncio +async def test_closing_disconnects_prisma_even_when_the_native_pool_fails_to_close(): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + native: Final = MagicMock() + native.close = AsyncMock(side_effect=RuntimeError("native close failed")) + + with pytest.raises(RuntimeError, match="native close failed"): + async with _db.lifespan(_db.Database(prisma=prisma, native=native)): + pass + native.close.assert_awaited_once() + prisma.disconnect.assert_awaited_once()