feat(db): add the SQLx pool foundation with the writer pool
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run

NativeDatabase.connect opens a lazy SQLx writer pool from an Endpoint that
Python resolves from the final DATABASE_URL, read the way Prisma reads it:
Unix socket host, sslmode / sslaccept / sslcert with Prisma's verification
rules (prefer + strict asks the server once whether it offers TLS),
sslidentity / sslpassword handed over as PEM, schema, options, pgbouncer mode,
pool_timeout (0 waits forever) and socket_timeout. Options SQLx cannot send
are logged instead of failing boot

ConnectionBudget splits each worker's connection_limit between Prisma and the
native pool, so the budget an operator sized never grows. The Rust share is the
hard-coded RUST_RESERVED_CONNECTIONS and no setting names Rust, SQLx or Prisma.
It is 0 until a ported query needs the pool, so Prisma's URL is exactly what it
was, no native pool opens and nothing connects through SQLx. A connection_limit
that is not a positive integer is never split and goes to Prisma as before.
Once the share is above 0 Prisma keeps at least one connection, and pools are
tagged litellm-prisma and litellm-native in pg_stat_activity. Python keeps
minting RDS IAM and Entra tokens and hands each one to the pool

statement_cache_size sets the native statement cache (Prisma's default 100) and
pgbouncer=true turns it off. sqlx-postgres 0.9.0 still names every statement when
the cache is off, so Statements runs queries unnamed in that case, like Prisma,
until a release carries transact-rs/sqlx#4337. An application_name the operator
set is used on both pools

The SQL query files are pinned to LF so Windows checkouts hash them like .sqlx,
and the crate lists .sqlx in its package include so the core sdist builds offline

crates/db/AGENTS.md tracks parity with the Python Prisma stack

The bridge now links litellm-db, so the editable install compiles its queries.
The Rust DB workflow builds that install from the checked-in .sqlx metadata,
since it runs before the migrations that the online check needs
This commit is contained in:
Yujong Lee 2026-11-13 11:04:45 +09:00
parent 4dad0bc513
commit dc097107b2
30 changed files with 1757 additions and 62 deletions

View file

@ -91,6 +91,8 @@ jobs:
- name: Install Python dependencies
working-directory: .
env:
SQLX_OFFLINE: "true"
run: uv sync --frozen --extra extra-proxy
- name: Deploy Prisma migrations

View file

@ -0,0 +1,52 @@
"""Split each Postgres pool's connection budget between Prisma and the Rust core's SQLx pool.
Migration scaffolding. Prisma and SQLx share a database only while the Python gateway still runs Prisma
queries next to ported ones. These invariants hold until Prisma is gone:
- Operators set one `connection_limit` per pool (`database_connection_pool_limit`, a URL pin or
`database_extra_connection_params`). Nothing an operator sets names Prisma, SQLx or Rust. The Rust share
is carved out of that limit, never added on top, so a Rust pool does not grow the connection budget.
- The Rust share is the hard-coded `RUST_RESERVED_CONNECTIONS`. It is 0 until a ported query runs on the
native pool, so Prisma keeps its whole limit, untouched, and no native pool opens.
- Only a positive integer limit is split. Any other value goes to Prisma as it is, which judges it the same
way it did before the split existed.
- The split is Python's. It is about Prisma and Prisma lives only in the Python gateway, so this module is
deleted with Prisma.
- The Rust core takes a plain `max_connections` and never learns about the split. That is already its final
interface, since the Rust gateway will pass `database_connection_pool_limit` straight to SQLx.
"""
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
RUST_RESERVED_CONNECTIONS: Final = 0
@dataclass(frozen=True, slots=True)
class ConnectionBudget:
total: int
python: int
rust: int
@classmethod
def split(cls, total: int, rust_reserved: int = RUST_RESERVED_CONNECTIONS) -> "ConnectionBudget":
"""Prisma keeps at least one connection, so a total at or below the reserve ends up over it."""
return cls(total=total, python=max(1, total - rust_reserved), rust=rust_reserved)
@classmethod
def of(
cls, connection_limit: str | float, rust_reserved: int = RUST_RESERVED_CONNECTIONS
) -> "ConnectionBudget | None":
"""The split of a positive integer limit, or None for any other value Prisma is left to judge."""
digits: Final = str(connection_limit)
if not (digits.isascii() and digits.isdigit()) or int(digits) == 0:
return None
return cls.split(int(digits), rust_reserved)
def prisma_params(self, params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]:
"""``params`` with Prisma's share as its ``connection_limit``, untouched while the Rust pool takes none."""
if self.rust == 0:
return params
return MappingProxyType({**params, "connection_limit": self.python})

View file

@ -254,12 +254,28 @@ def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = p
return urllib.parse.urlunsplit(parsed._replace(query=query))
def postgres_connection_budget_message(writer_limit: str, reader_limit: str | None, num_workers: str) -> str:
PRISMA_APPLICATION_NAME: Final = "litellm-prisma"
NATIVE_APPLICATION_NAME: Final = "litellm-native"
def prisma_default_params(
lifetime_params: Mapping[str, str | int | float], native_connections: int
) -> Mapping[str, str | int | float]:
"""Prisma's URL defaults, tagged ``PRISMA_APPLICATION_NAME`` only once a native pool shares the database."""
if native_connections == 0:
return lifetime_params
return MappingProxyType({**lifetime_params, "application_name": PRISMA_APPLICATION_NAME})
def postgres_connection_budget_message(
writer_limit: str, rust_writer_limit: int, reader_limit: str | None, num_workers: str
) -> str:
"""The startup line that states this pod's worst-case Postgres connection demand.
Prisma's ``connection_limit`` is per query engine, and every uvicorn worker
owns one engine per configured database (writer, plus the reader when
``DATABASE_URL_READ_REPLICA`` is set). The limits are read off the final URLs,
``DATABASE_URL_READ_REPLICA`` is set) plus the Rust writer pool's share
(``connection_budget``) when it is above 0. The limits are read off the final URLs,
so a reader that pins its own ``connection_limit`` or an operator override in
``database_extra_connection_params`` is counted at its real value. The
server-side cap is shared by every pod, so the number an operator has to keep
@ -272,16 +288,18 @@ def postgres_connection_budget_message(writer_limit: str, reader_limit: str | No
reader: Final = int(reader_limit) if reader_limit is not None else 0
except ValueError:
return (
"LiteLLM Proxy: Postgres connection budget per pod = workers x (writer connection_limit + reader "
f"connection_limit) (workers={num_workers!r}, writer={writer_limit!r}, reader={reader_limit!r}); "
"LiteLLM Proxy: Postgres connection budget per pod = workers x (writer connection_limit + "
f"{'rust writer pool + ' if rust_writer_limit else ''}reader connection_limit) "
f"(workers={num_workers!r}, writer={writer_limit!r}, reader={reader_limit!r}); "
"keep pods x that figure under max_connections minus superuser_reserved_connections"
)
engines: Final = (
f"(writer connection_limit {writer} + reader connection_limit {reader})"
if reader_limit is not None
else f"writer connection_limit {writer}"
terms: Final = (
f"writer connection_limit {writer}",
*((f"rust writer pool {rust_writer_limit}",) if rust_writer_limit else ()),
*((f"reader connection_limit {reader}",) if reader_limit is not None else ()),
)
per_pod: Final = workers * (writer + reader)
engines: Final = terms[0] if len(terms) == 1 else f"({' + '.join(terms)})"
per_pod: Final = workers * (writer + rust_writer_limit + reader)
return (
f"LiteLLM Proxy: Postgres connection budget per pod = {workers} worker(s) x {engines} = up to {per_pod} "
"connections; keep pods x that figure under max_connections minus superuser_reserved_connections"

View file

@ -0,0 +1,281 @@
"""Settings for the native (SQLx) pool, resolved from the same URL Prisma connects with.
``DATABASE_URL`` is final once ``db_url_settings`` and the CLI have assembled it: token
minted, libpq TLS translated to Prisma's dialect, timeouts in ``options``, pool params
added. Reading that URL keeps one resolver for both pools instead of a second pass over
the ``DATABASE_*`` env vars.
Credentials follow the host. Each process runs one token refresh loop, owned by whatever
hosts it, and feeds every client in that process: here the worker's loop re-mints Prisma's
URL and pushes the same kind of token into the Rust pool through ``set_writer_password``.
The Rust core only accepts a password and never mints one, so a Rust gateway process
brings its own loop. Two processes minting at once is fine, since IAM and Entra tokens
are independent. What must not exist twice is the logic that resolves region, endpoint
and user, so when the Rust gateway needs it, it moves into the Rust core and this
loop calls it.
"""
import asyncio
import shlex
import struct
import urllib.parse
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final, Protocol
from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat, pkcs12
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.db_url_settings import NATIVE_APPLICATION_NAME, PRISMA_APPLICATION_NAME
from litellm.proxy.db.pgbouncer import database_url_is_pooled, fold_libpq_option_tokens, libpq_option_setting
from litellm.proxy.db.token_auth import (
DEFAULT_POSTGRES_PORT,
DatabaseTokenAuth,
IAMEndpoint,
mint_database_token,
parse_database_token_expiration,
parse_iam_endpoint_from_url,
resolve_database_token_auth,
)
from litellm.rust_bridge.database import (
NativeClientIdentity,
NativeDatabaseEndpoint,
NativePoolLimits,
SslMode,
)
TOKEN_REFRESH_BUFFER_SECONDS: Final = 180.0
TOKEN_FALLBACK_REFRESH_SECONDS: Final = 600.0
TOKEN_RETRY_SECONDS: Final = 30.0
PRISMA_DEFAULT_SCHEMA: Final = "public"
PRISMA_DEFAULT_POOL_TIMEOUT_SECONDS: Final = 10.0
PRISMA_DEFAULT_STATEMENT_CACHE_SIZE: Final = 100
TLS_PROBE_TIMEOUT_SECONDS: Final = 5.0
_POSTGRES_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103)
@dataclass(frozen=True, slots=True)
class NativeEndpointSettings:
endpoint: NativeDatabaseEndpoint
ignored_options: tuple[str, ...]
def _params(url: str) -> Mapping[str, str]:
return MappingProxyType(dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query)))
def needs_tls_probe(url: str) -> bool:
"""Prisma's ``sslmode=prefer`` with ``sslaccept=strict`` verifies TLS when the server offers it
and falls back to plaintext when it does not; SQLx has no verified-or-plaintext mode, so the
caller asks the server once and pins the answer."""
params: Final = _params(url)
return params.get("sslmode", "prefer") == "prefer" and params.get("sslaccept") == "strict" and "host" not in params
async def server_offers_tls(url: str, timeout: float = TLS_PROBE_TIMEOUT_SECONDS) -> bool:
"""Whether the server answers a Postgres ``SSLRequest`` with ``S``; an unreachable server counts as offering it."""
parsed: Final = urllib.parse.urlsplit(url)
try:
reader, writer = await asyncio.wait_for(
asyncio.open_connection(parsed.hostname, parsed.port or int(DEFAULT_POSTGRES_PORT)), timeout
)
except (OSError, asyncio.TimeoutError):
return True
try:
writer.write(_POSTGRES_SSL_REQUEST)
await writer.drain()
answer: Final = await asyncio.wait_for(reader.read(1), timeout)
except (OSError, asyncio.TimeoutError):
return True
finally:
writer.close()
return answer != b"N"
def _tls(params: Mapping[str, str], tls_offered: bool) -> tuple[SslMode, str | None]:
sslmode: Final = params.get("sslmode", "prefer")
strict: Final = params.get("sslaccept") == "strict"
if sslmode == "disable" or (sslmode == "prefer" and strict and not tls_offered):
return "disable", None
if strict:
return "verify-full", params.get("sslcert")
return ("require" if sslmode == "require" else "prefer"), None
def _client_identity(
params: Mapping[str, str], ssl_mode: SslMode, read_file: Callable[[Path], bytes]
) -> NativeClientIdentity | None:
identity_path: Final = params.get("sslidentity")
if identity_path is None or ssl_mode == "disable":
return None
password: Final = params.get("sslpassword")
key, cert, chain = pkcs12.load_key_and_certificates(
read_file(Path(identity_path)), None if password is None else password.encode()
)
if key is None or cert is None:
raise ValueError(f"sslidentity {identity_path!r} must hold a client certificate and its private key")
return NativeClientIdentity(
cert_pem=b"".join(c.public_bytes(Encoding.PEM) for c in (cert, *chain)).decode(),
key_pem=key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption()).decode(),
)
def _session_options(params: Mapping[str, str]) -> tuple[tuple[tuple[str, str], ...], tuple[str, ...]]:
raw: Final = params.get("options", "")
try:
tokens: Final[Sequence[str]] = fold_libpq_option_tokens(tuple(shlex.split(raw)))
except ValueError:
return (), (raw,)
readable: Final = tuple((token, libpq_option_setting(token)) for token in tokens)
schema: Final = params.get("schema", PRISMA_DEFAULT_SCHEMA)
search_path: Final = () if schema == PRISMA_DEFAULT_SCHEMA else (("search_path", f'"{schema}"'),)
settings: Final = tuple(
(name, value) for name, _, value in (setting.partition("=") for _, setting in readable if setting is not None)
)
return (*search_path, *settings), tuple(token for token, setting in readable if setting is None)
def _seconds(params: Mapping[str, str], key: str, default: float | None = None) -> float | None:
value: Final = params.get(key)
if value is None:
return default
return None if float(value) == 0 else float(value)
def _statement_cache_capacity(params: Mapping[str, str]) -> int:
if params.get("pgbouncer") == "true":
return 0
return int(params.get("statement_cache_size", PRISMA_DEFAULT_STATEMENT_CACHE_SIZE))
def _application_name(params: Mapping[str, str]) -> str:
configured: Final = params.get("application_name", PRISMA_APPLICATION_NAME)
return NATIVE_APPLICATION_NAME if configured == PRISMA_APPLICATION_NAME else configured
def native_endpoint(
url: str,
*,
tls_offered: bool,
max_connections: int,
read_file: Callable[[Path], bytes] = Path.read_bytes,
) -> NativeEndpointSettings:
"""The native pool's endpoint for the Prisma URL ``url``, read the way Prisma reads it.
``host`` is Prisma's Unix socket directory. The pool takes ``max_connections``, the
Rust share of the budget (``connection_budget``), and tags its connections
``NATIVE_APPLICATION_NAME``. An unset ``pool_timeout`` takes Prisma's default and
0 waits forever.
``sslidentity`` / ``sslpassword`` are a PKCS#12 client certificate, handed over as PEM.
``options`` tokens that are not ``-c name=value`` settings cannot be sent by SQLx and
come back in ``ignored_options`` instead of failing boot.
"""
parsed: Final = urllib.parse.urlsplit(url)
params: Final = _params(url)
socket: Final = params.get("host")
if not (parsed.hostname or socket) or not parsed.username:
raise ValueError("DATABASE_URL must name a host and a user for the native pool")
ssl_mode, ssl_root_cert = _tls(params, tls_offered)
options, ignored = _session_options(params)
return NativeEndpointSettings(
endpoint=NativeDatabaseEndpoint(
host=parsed.hostname or "localhost",
port=parsed.port or int(DEFAULT_POSTGRES_PORT),
socket=socket,
username=urllib.parse.unquote(parsed.username),
password=urllib.parse.unquote(parsed.password) if parsed.password else None,
database=urllib.parse.unquote(parsed.path.lstrip("/")) or "postgres",
ssl_mode=ssl_mode,
ssl_root_cert=ssl_root_cert,
ssl_client=_client_identity(params, ssl_mode, read_file),
application_name=_application_name(params),
options=options,
statement_cache_capacity=_statement_cache_capacity(params),
pool=NativePoolLimits(
max_connections=max_connections,
acquire_timeout_seconds=_seconds(params, "pool_timeout", PRISMA_DEFAULT_POOL_TIMEOUT_SECONDS),
idle_timeout_seconds=_seconds(params, "max_idle_connection_lifetime"),
socket_timeout_seconds=_seconds(params, "socket_timeout"),
),
),
ignored_options=ignored,
)
async def resolve_native_endpoint(url: str, max_connections: int) -> NativeEndpointSettings:
tls_offered: Final = await server_offers_tls(url) if needs_tls_probe(url) else True
return native_endpoint(url, tls_offered=tls_offered, max_connections=max_connections)
class PasswordSink(Protocol):
def set_writer_password(self, password: str) -> None: ...
@dataclass(frozen=True, slots=True)
class TokenRotation:
auth: DatabaseTokenAuth
endpoint: IAMEndpoint
mint_token: Callable[[DatabaseTokenAuth, IAMEndpoint], str] = mint_database_token
def mint(self) -> str:
return urllib.parse.unquote(self.mint_token(self.auth, self.endpoint))
def seconds_until_refresh(self, token: str, now: datetime) -> float:
expires_at: Final = parse_database_token_expiration(self.auth, token)
if expires_at is None:
return TOKEN_FALLBACK_REFRESH_SECONDS
return max(TOKEN_RETRY_SECONDS, (expires_at - now).total_seconds() - TOKEN_REFRESH_BUFFER_SECONDS)
def writer_token_rotation(url: str) -> TokenRotation | None:
"""The token rotation the writer needs, or None for a static password.
A URL pointed at the in-container PgBouncer carries PgBouncer's own credential,
and PgBouncer renews the upstream token itself.
"""
if database_url_is_pooled():
return None
auth: Final = resolve_database_token_auth()
if auth is None:
return None
return TokenRotation(auth=auth, endpoint=parse_iam_endpoint_from_url(url))
def _utcnow() -> datetime:
return datetime.now(timezone.utc).replace(tzinfo=None)
async def _refresh(sink: PasswordSink, rotation: TokenRotation, now: Callable[[], datetime]) -> float:
try:
token: Final = await asyncio.to_thread(rotation.mint)
except Exception as error: # noqa: BLE001 # boto3 and azure-identity raise many types; any failure means retry
verbose_proxy_logger.error(
"Native database pool keeps its current %s (%r); retrying in %.0fs.",
rotation.auth.label,
error,
TOKEN_RETRY_SECONDS,
)
return TOKEN_RETRY_SECONDS
sink.set_writer_password(token)
return rotation.seconds_until_refresh(token, now())
async def keep_writer_token_fresh(
sink: PasswordSink,
rotation: TokenRotation,
current_token: str,
now: Callable[[], datetime] = _utcnow,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
) -> None:
"""Hand the native pool a fresh token before the current one expires.
Existing connections stay authenticated; only connections opened after a refresh
use the new token, the same as a re-minted Prisma URL.
"""
delay = rotation.seconds_until_refresh(current_token, now()) # rebind-ok: each refresh schedules the next
while True:
await sleep(delay)
delay = await _refresh(sink, rotation, now)

View file

@ -140,19 +140,27 @@ def _userlist_quote(value: str) -> str:
return '"' + value.replace('"', '""') + '"'
def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None:
"""The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else.
Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a
detached ``-c`` is folded into the token that follows it first.
"""
folded: Final = tuple(
def fold_libpq_option_tokens(tokens: Sequence[str]) -> tuple[str, ...]:
"""The tokens of a libpq ``options`` string with each detached ``-c`` folded into the token after it."""
return tuple(
f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token
for index, token in enumerate(tokens)
if index == 0 or tokens[index - 1] != "-c"
)
settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:])
return settings if len(settings) == len(folded) else None
def libpq_option_setting(token: str) -> str | None:
"""The ``name=value`` setting a folded ``-cname=value`` or ``--name=value`` token carries, or None."""
return token[2:] if token.startswith(("-c", "--")) and "=" in token[2:] else None
def libpq_option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None:
"""The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else.
Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``.
"""
settings: Final = tuple(libpq_option_setting(token) for token in fold_libpq_option_tokens(tokens))
return None if None in settings else tuple(setting for setting in settings if setting is not None)
def _connect_query(options: str) -> str | PgBouncerError:
@ -163,7 +171,7 @@ def _connect_query(options: str) -> str | PgBouncerError:
the settings are applied to each new server connection instead. Every
client shares them, which is what the single ``DATABASE_URL`` gave anyway.
"""
settings: Final = _option_settings(tuple(shlex.split(options)))
settings: Final = libpq_option_settings(tuple(shlex.split(options)))
if settings is None:
return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings")
return "; ".join(

View file

@ -79,7 +79,7 @@ def _build_db_connection_url_params(
socket_timeout: float | None = None,
disable_prepared_statements: bool = False,
extra_params: dict | None = None,
) -> dict:
) -> dict[str, str | int | float]:
"""Build the Prisma DATABASE_URL query params controlling connection pool behavior.
`connect_timeout` / `socket_timeout` map to the Prisma URL params of the same
@ -91,7 +91,7 @@ def _build_db_connection_url_params(
`extra_params` is an untyped passthrough — keys it provides win over the
named arguments above, so it can be used to override any default we set here.
"""
params: Final[dict] = {
params: Final[dict[str, str | int | float]] = {
"connection_limit": connection_limit,
}
if pool_timeout is not None:
@ -1637,11 +1637,13 @@ def run_server(
db_connection_timeout = LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value
if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None:
from litellm.proxy.db.connection_budget import ConnectionBudget
from litellm.proxy.db.db_url_settings import (
DISABLE_PREPARED_STATEMENTS_ENV_VAR,
add_missing_query_params,
idle_lifetime_params,
postgres_connection_budget_message,
prisma_default_params,
reader_shareable_params,
translate_libpq_ssl_params,
unsupported_db_scheme,
@ -1676,7 +1678,13 @@ def run_server(
disable_prepared_statements=disable_prepared_statements,
extra_params=db_extra_connection_params,
)
writer_budget: Final = ConnectionBudget.of(connection_url_params["connection_limit"])
rust_writer_connections: Final = 0 if writer_budget is None else writer_budget.rust
prisma_url_params: Final = (
connection_url_params if writer_budget is None else writer_budget.prisma_params(connection_url_params)
)
lifetime_params: Final = idle_lifetime_params(general_settings.get("database_max_idle_connection_lifetime"))
prisma_defaults: Final = prisma_default_params(lifetime_params, rust_writer_connections)
if os.getenv("DATABASE_URL", None) is not None:
database_url = get_secret("DATABASE_URL", default_value=None)
resolved_url: Final[str | None] = str(database_url) if database_url else None
@ -1693,10 +1701,10 @@ def run_server(
)
modified_url = append_query_params(
writer_url,
connection_url_params,
dict(prisma_url_params),
)
os.environ["DATABASE_URL"] = translate_libpq_ssl_params(
add_missing_query_params(modified_url, lifetime_params)
add_missing_query_params(modified_url, prisma_defaults)
)
if os.getenv("DIRECT_URL", None) is not None:
database_url = os.getenv("DIRECT_URL")
@ -1723,15 +1731,16 @@ def run_server(
_with_query_value(read_replica_url, "options", reader_options)
if reader_options
else read_replica_url,
reader_shareable_params(connection_url_params),
reader_shareable_params(prisma_url_params),
),
lifetime_params,
prisma_defaults,
)
)
print_startup_message(
postgres_connection_budget_message(
writer_limit=_url_query_value(os.getenv("DATABASE_URL"), "connection_limit")
or f"{db_connection_pool_limit}",
rust_writer_limit=rust_writer_connections,
reader_limit=(
_url_query_value(os.getenv("DATABASE_URL_READ_REPLICA"), "connection_limit")
or f"{db_connection_pool_limit}"

View file

@ -667,6 +667,7 @@ from litellm.proxy.config_resolvers.settings_rules import (
from litellm.proxy.config_resolvers.sso import SSO_FIELD_ENV_VARS
from litellm.proxy.container_endpoints.endpoints import router as container_router
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.connection_budget import RUST_RESERVED_CONNECTIONS
from litellm.proxy.db.create_views import SupportsRawQueries
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
@ -1601,7 +1602,9 @@ async def _connect_to_count_stored_values(resources: AsyncExitStack) -> Supports
async def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None:
if client is None:
return None
database: Final = await database_runtime.Database.open(client)
database: Final = await database_runtime.Database.open(
client, os.environ["DATABASE_URL"], RUST_RESERVED_CONNECTIONS
)
resources.push_async_callback(database.aclose)
return database

View file

@ -1,9 +1,16 @@
import asyncio
import urllib.parse
from collections.abc import AsyncGenerator, Coroutine
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.native_database import (
keep_writer_token_fresh,
resolve_native_endpoint,
writer_token_rotation,
)
from litellm.rust_bridge._native import NativeDatabase
if TYPE_CHECKING:
@ -13,22 +20,47 @@ if TYPE_CHECKING:
@dataclass(frozen=True, slots=True)
class Database:
prisma: "PrismaClient"
native: NativeDatabase
native: NativeDatabase | None
tasks: "tuple[asyncio.Task[None], ...]" = ()
@classmethod
async def open(cls, prisma: "PrismaClient") -> "Database":
async def open(cls, prisma: "PrismaClient", writer_url: str, native_connections: int) -> "Database":
if native_connections == 0:
return cls(prisma=prisma, native=None)
try:
native: Final = await NativeDatabase.connect()
settings: Final = await resolve_native_endpoint(writer_url, native_connections)
native: Final = await NativeDatabase.connect(settings.endpoint)
except BaseException:
await prisma.aclose()
raise
return cls(prisma=prisma, native=native)
if settings.ignored_options:
verbose_proxy_logger.error(
"Native database pool cannot send the DATABASE_URL options %s; only -c name=value settings carry over",
settings.ignored_options,
)
rotation: Final = writer_token_rotation(writer_url)
token_refresh: Final = (
()
if rotation is None
else (
asyncio.create_task(
keep_writer_token_fresh(
native, rotation, urllib.parse.unquote(urllib.parse.urlsplit(writer_url).password or "")
)
),
)
)
return cls(prisma=prisma, native=native, tasks=token_refresh)
async def aclose(self) -> None:
for task in self.tasks:
task.cancel()
await asyncio.gather(*self.tasks, return_exceptions=True)
try:
await self.prisma.aclose()
finally:
await self.native.close()
if self.native is not None:
await self.native.close()
@asynccontextmanager

View file

@ -3766,6 +3766,8 @@ name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"futures-util",
"litellm-secret-value",
"macro_rules_attribute",
"rstest",
"serde",
@ -3774,6 +3776,7 @@ dependencies = [
"strum",
"thiserror 2.0.19",
"tokio",
"url",
]
[[package]]
@ -4405,6 +4408,7 @@ dependencies = [
"litellm-callbacks-legacy-python",
"litellm-circuit-breaker",
"litellm-core-utils",
"litellm-db",
"litellm-error",
"litellm-host",
"litellm-host-python",

View file

@ -23,3 +23,38 @@ SQLX_OFFLINE=false cargo sqlx prepare
```
Tests use transactions and leave them uncommitted so fixture rows roll back when each test finishes
# Boundary with the hosts
The Python gateway and the Rust gateway both host this crate during the migration, so it takes only what both can hand it:
- Pool size is a plain `max_connections`. How a worker's `connection_limit` is split between Prisma and SQLx is decided by the Python host (`litellm/proxy/db/connection_budget.py`) and disappears with Prisma. Never add budget or Prisma knowledge here
- Credentials arrive as a password (`set_writer_password`). Each process runs one token refresh loop owned by its host, so this crate never mints or refreshes a token on its own. When the Rust gateway needs IAM or Entra tokens, the minting logic moves here as a function the host's loop calls, and the Python loop switches to it, so the resolution logic exists once
# Parity with the Python Prisma stack
Python and SQLx run side by side until every repository operation is native, so the SQLx pool must honor every database setting Python honors. `Endpoint` is resolved in Python from the final `DATABASE_URL` (`litellm/proxy/db/native_database.py`); Rust never reads `DATABASE_*` env vars or parses a Prisma URL. When `litellm/proxy/db/` gains a setting or behavior, add it here
| Python behavior | Where | Native status |
| --- | --- | --- |
| URL assembled from `DATABASE_*`, token minted, libpq TLS translated | `db_url_settings.py` | Inherited: the native endpoint reads the final URL |
| `host` Unix socket directory | URL param | `Endpoint::socket` |
| `sslmode`, `sslaccept=strict`, `sslcert` | `db_url_settings.py` | `strict` is `verify-full` with `sslcert` as the root; without it nothing is verified, so `sslcert` is dropped. `prefer` + `strict` asks the server once at startup and pins `verify-full` or `disable` |
| `sslidentity`, `sslpassword` | URL params | Python opens the PKCS#12 file and hands the client certificate and key over as PEM |
| `connection_limit` | `proxy_cli.py`, `connection_budget.py` | Split, not added: `ConnectionBudget.split` gives the Rust pool the hard-coded `RUST_RESERVED_CONNECTIONS` and Prisma the rest, so an upgrade never grows the budget. No setting names Rust, SQLx or Prisma. The reserve is 0 until a ported query needs the pool, so Prisma keeps its whole limit and no native pool opens. Once it is above 0, Prisma connections are tagged `application_name=litellm-prisma` and native ones `litellm-native` |
| `pool_timeout`, `max_idle_connection_lifetime` | `proxy_cli.py`, `db_url_settings.py` | `PoolLimits`, with Prisma's 10s default. `pool_timeout=0` waits forever |
| `socket_timeout` | `proxy_cli.py` | `PoolLimits::socket_timeout_seconds`, a deadline on each native operation |
| `connect_timeout` | `proxy_cli.py` | Not carried over: SQLx has no per-connect timeout, so a black-holed connect fails at `pool_timeout` |
| `options` (`statement_timeout`, `lock_timeout`, `idle_in_transaction_session_timeout`) | `proxy_cli.py` | `Endpoint::options`, sent as startup `-c` settings. Tokens that are not `-c name=value` are logged at boot and dropped |
| `schema` | URL param | `search_path` startup setting when not `public`; a custom schema behind PgBouncer is rejected by PgBouncer |
| `statement_cache_size`, `pgbouncer=true`, `DATABASE_DISABLE_PREPARED_STATEMENTS` | URL params, `db_url_settings.py` | `statement_cache_capacity`: `statement_cache_size` (Prisma's default 100), 0 under `pgbouncer=true`. Run every query through `Statements::new(executor, endpoint.names_statements())` so a 0 capacity sends unnamed statements like Prisma, until a sqlx release carries transact-rs/sqlx#4337 |
| `application_name` | URL param | A value the operator set is used on both pools; Prisma's `litellm-prisma` default becomes `litellm-native` |
| Rotating RDS IAM / Entra token | `token_auth.py`, `PrismaWrapper` | Python mints, `NativeDatabase.set_writer_password` swaps connect options |
| In-container PgBouncer | `pgbouncer.py` | Inherited: the pooled URL carries PgBouncer's credential |
| Startup connection budget | `postgres_connection_budget_message` | Counts the Rust writer pool's share |
| Read replica routing | `routing_prisma_wrapper.py` | Planned: optional reader pool with explicit routing |
| Error classification, `lock_timeout` on rollups | `exception_handler.py`, `rollup_lock_timeout.py` | Planned |
| `ServiceTypes.DB` events per round trip | `db_span.py`, `log_db_metrics.py` | Planned |
| Auth lookup gate, fail open on DB outage, reader choice per query | `db_lookup_gate.py`, `allow_requests_on_db_unavailable` | With each ported auth read |
| Health watchdog, engine reaper, drain tracker, engine recreation | `utils.py`, `query_engine_reaper.py`, `prisma_client.py` | Not needed: SQLx pools reconnect per connection and `close` waits for checked-out connections |
| Migrations, `check_migration`, views | `prisma_client.py`, `create_views.py` | Prisma keeps them until the SQLx cutover |

View file

@ -4,16 +4,21 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
include = ["Cargo.toml", "src/**", "queries/**", ".sqlx/**"]
[dependencies]
chrono.workspace = true
futures-util.workspace = true
litellm-secret-value.workspace = true
macro_rules_attribute.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
strum.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["time"] }
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true
tokio = { workspace = true, features = ["test-util"] }
url.workspace = true

View file

@ -0,0 +1 @@
*.sql text eol=lf

View file

@ -0,0 +1,51 @@
use std::{io, time::Duration};
use crate::Error;
pub async fn within<T>(
socket_timeout: Option<Duration>,
operation: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let Some(limit) = socket_timeout else {
return operation.await;
};
tokio::time::timeout(limit, operation)
.await
.unwrap_or_else(|_| {
Err(Error::Query(sqlx::Error::Io(
io::ErrorKind::TimedOut.into(),
)))
})
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::within;
use crate::Error;
#[tokio::test(start_paused = true)]
async fn a_slow_operation_fails_like_a_socket_timeout() {
let result = within(Some(Duration::from_secs(5)), async {
tokio::time::sleep(Duration::from_secs(6)).await;
Ok::<_, Error>(())
})
.await;
assert!(
matches!(result, Err(Error::Query(sqlx::Error::Io(error))) if error.kind() == std::io::ErrorKind::TimedOut)
);
}
#[tokio::test(start_paused = true)]
async fn no_timeout_waits_for_the_operation() {
let result = within(None, async {
tokio::time::sleep(Duration::from_secs(3600)).await;
Ok::<_, Error>(7)
})
.await;
assert_eq!(result.ok(), Some(7));
}
}

View file

@ -0,0 +1,332 @@
use std::{path::PathBuf, time::Duration};
use litellm_secret_value::SecretValue;
use serde::Deserialize;
use sqlx::{
PgPool,
postgres::{PgConnectOptions, PgPoolOptions, PgSslMode},
};
use crate::Error;
const WAIT_FOREVER: Duration = Duration::from_secs(100 * 365 * 24 * 60 * 60);
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub struct Endpoint {
pub host: String,
pub port: u16,
pub socket: Option<PathBuf>,
pub username: String,
pub password: Option<SecretValue>,
pub database: String,
pub ssl_mode: SslMode,
pub ssl_root_cert: Option<PathBuf>,
pub ssl_client: Option<ClientIdentity>,
pub application_name: Option<String>,
pub options: Vec<(String, String)>,
pub statement_cache_capacity: usize,
pub pool: PoolLimits,
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub enum SslMode {
Disable,
Prefer,
Require,
VerifyCa,
VerifyFull,
}
#[derive(Clone, Debug, Deserialize, PartialEq)]
pub struct ClientIdentity {
pub cert_pem: String,
pub key_pem: SecretValue,
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub struct PoolLimits {
pub max_connections: u32,
pub acquire_timeout_seconds: Option<f64>,
pub idle_timeout_seconds: Option<f64>,
pub socket_timeout_seconds: Option<f64>,
}
impl From<SslMode> for PgSslMode {
fn from(mode: SslMode) -> Self {
match mode {
SslMode::Disable => Self::Disable,
SslMode::Prefer => Self::Prefer,
SslMode::Require => Self::Require,
SslMode::VerifyCa => Self::VerifyCa,
SslMode::VerifyFull => Self::VerifyFull,
}
}
}
impl Endpoint {
pub fn with_password(&self, password: SecretValue) -> Self {
Self {
password: Some(password),
..self.clone()
}
}
pub fn connect_options(&self) -> PgConnectOptions {
let options = PgConnectOptions::new_without_pgpass()
.host(&self.host)
.port(self.port)
.username(&self.username)
.database(&self.database)
.ssl_mode(self.ssl_mode.into())
.options(self.options.iter().map(|(name, value)| (name, value)))
.statement_cache_capacity(self.statement_cache_capacity);
let options = match &self.socket {
Some(directory) => options.socket(directory),
None => options,
};
let options = match &self.password {
Some(password) => options.password(password.expose()),
None => options,
};
let options = match &self.application_name {
Some(name) => options.application_name(name),
None => options,
};
let options = match &self.ssl_client {
Some(identity) => options
.ssl_client_cert_from_pem(&identity.cert_pem)
.ssl_client_key_from_pem(identity.key_pem.expose()),
None => options,
};
match &self.ssl_root_cert {
Some(cert) => options.ssl_root_cert(cert),
None => options,
}
}
pub fn pool(&self) -> Result<PgPool, Error> {
let limits = self.pool;
Ok(PgPoolOptions::new()
.min_connections(0)
.max_connections(limits.max_connections)
.acquire_timeout(seconds(limits.acquire_timeout_seconds)?.unwrap_or(WAIT_FOREVER))
.idle_timeout(seconds(limits.idle_timeout_seconds)?)
.connect_lazy_with(self.connect_options()))
}
pub fn names_statements(&self) -> bool {
self.statement_cache_capacity > 0
}
pub fn socket_timeout(&self) -> Result<Option<Duration>, Error> {
seconds(self.pool.socket_timeout_seconds)
}
}
fn seconds(value: Option<f64>) -> Result<Option<Duration>, Error> {
value
.map(Duration::try_from_secs_f64)
.transpose()
.map_err(Error::from)
}
#[cfg(test)]
mod tests {
use std::{error::Error as StdError, path::PathBuf, time::Duration};
use litellm_secret_value::SecretValue;
use rstest::rstest;
use sqlx::postgres::PgSslMode;
use super::{Endpoint, PoolLimits, SslMode, WAIT_FOREVER};
use crate::{Error, test_support::database_endpoint};
const LIMITS: PoolLimits = PoolLimits {
max_connections: 7,
acquire_timeout_seconds: Some(2.5),
idle_timeout_seconds: Some(60.0),
socket_timeout_seconds: None,
};
fn endpoint() -> Endpoint {
Endpoint {
host: "db.internal".to_owned(),
port: 6543,
socket: None,
username: "litellm".to_owned(),
password: Some(SecretValue::new("first")),
database: "proxy".to_owned(),
ssl_mode: SslMode::VerifyFull,
ssl_root_cert: None,
ssl_client: None,
application_name: Some("litellm-native".to_owned()),
options: vec![("statement_timeout".to_owned(), "30000".to_owned())],
statement_cache_capacity: 100,
pool: LIMITS,
}
}
#[test]
fn maps_endpoint_to_connect_options() {
let options = endpoint().connect_options();
assert_eq!(options.get_host(), "db.internal");
assert_eq!(options.get_port(), 6543);
assert_eq!(options.get_socket(), None);
assert_eq!(options.get_username(), "litellm");
assert_eq!(options.get_database(), Some("proxy"));
assert_eq!(options.get_application_name(), Some("litellm-native"));
assert!(matches!(options.get_ssl_mode(), PgSslMode::VerifyFull));
assert_eq!(options.get_options(), Some("-c statement_timeout=30000"));
}
#[test]
fn connects_through_a_unix_socket_directory() {
let options = Endpoint {
socket: Some(PathBuf::from("/cloudsql/project:region:instance")),
..endpoint()
}
.connect_options();
assert_eq!(
options.get_socket(),
Some(&PathBuf::from("/cloudsql/project:region:instance"))
);
}
#[tokio::test]
async fn applies_pool_limits() -> Result<(), Error> {
let pool = endpoint().pool()?;
assert_eq!(pool.options().get_max_connections(), 7);
assert_eq!(pool.options().get_min_connections(), 0);
assert_eq!(
pool.options().get_acquire_timeout(),
Duration::from_millis(2500)
);
assert_eq!(
pool.options().get_idle_timeout(),
Some(Duration::from_secs(60))
);
assert_eq!(pool.size(), 0);
Ok(())
}
#[tokio::test]
async fn missing_acquire_timeout_waits_forever() -> Result<(), Error> {
let pool = Endpoint {
pool: PoolLimits {
acquire_timeout_seconds: None,
..LIMITS
},
..endpoint()
}
.pool()?;
assert_eq!(pool.options().get_acquire_timeout(), WAIT_FOREVER);
Ok(())
}
#[rstest]
#[case::negative(-1.0)]
#[case::not_a_number(f64::NAN)]
#[tokio::test]
async fn rejects_invalid_durations(#[case] seconds: f64) {
let error = Endpoint {
pool: PoolLimits {
acquire_timeout_seconds: Some(seconds),
..LIMITS
},
..endpoint()
}
.pool()
.expect_err("an invalid duration must not build a pool");
assert!(matches!(error, Error::InvalidDuration(_)));
}
#[test]
fn redacts_password() {
assert!(!format!("{:?}", endpoint()).contains("first"));
}
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn waits_past_an_exhausted_pool_without_a_timeout() -> Result<(), Box<dyn StdError>> {
let endpoint = database_endpoint()?;
let pool = Endpoint {
pool: PoolLimits {
max_connections: 1,
acquire_timeout_seconds: None,
..endpoint.pool
},
..endpoint
}
.pool()?;
let held = pool.acquire().await?;
let waiter = tokio::spawn({
let pool = pool.clone();
async move { pool.acquire().await.map(drop) }
});
tokio::time::sleep(Duration::from_millis(300)).await;
assert!(!waiter.is_finished());
drop(held);
waiter.await??;
Ok(())
}
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn applies_session_options() -> Result<(), Box<dyn StdError>> {
let pool = Endpoint {
options: vec![
("statement_timeout".to_owned(), "1234".to_owned()),
(
"search_path".to_owned(),
"\"Other Schema\", public".to_owned(),
),
],
application_name: Some("litellm-native".to_owned()),
..database_endpoint()?
}
.pool()?;
#[expect(
clippy::disallowed_methods,
reason = "Session settings are not part of the compile-checked schema"
)]
let (statement_timeout, search_path, application_name): (String, String, String) =
sqlx::query_as(
"SELECT current_setting('statement_timeout'), current_setting('search_path'), \
current_setting('application_name')",
)
.fetch_one(&pool)
.await?;
assert_eq!(statement_timeout, "1234ms");
assert_eq!(search_path, "\"Other Schema\", public");
assert_eq!(application_name, "litellm-native");
Ok(())
}
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn rotated_password_reaches_new_connections() -> Result<(), Box<dyn StdError>> {
let endpoint = database_endpoint()?;
let pool = endpoint
.with_password(SecretValue::new("not-the-password"))
.pool()?;
#[expect(clippy::disallowed_methods, reason = "A probe query needs no schema")]
let probe = || sqlx::query("SELECT 1");
let stale = probe().execute(&pool).await;
pool.set_connect_options(endpoint.connect_options());
probe().execute(&pool).await?;
assert!(stale.is_err());
Ok(())
}
}

View file

@ -6,6 +6,8 @@ pub enum Error {
InvalidSpendBatch { ids: usize, costs: usize },
#[error("invalid spend timestamp")]
InvalidTimestamp(#[from] chrono::ParseError),
#[error("invalid pool duration")]
InvalidDuration(#[from] std::time::TryFromFloatSecsError),
}
impl Error {

View file

@ -7,15 +7,20 @@ macro_rules_attribute::attribute_alias! {
}
mod daily_spend;
mod deadline;
mod end_user;
mod endpoint;
mod error;
mod failure;
mod spend;
mod statements;
#[cfg(test)]
mod test_support;
pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
pub use deadline::within;
pub use end_user::{EndUser, RateLimits, find_end_user};
pub use endpoint::{ClientIdentity, Endpoint, PoolLimits, SslMode};
pub use error::Error;
pub use failure::DbFailure;
pub use spend::{
@ -24,3 +29,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 statements::Statements;

View file

@ -0,0 +1,152 @@
use futures_util::{future::BoxFuture, stream::BoxStream};
use sqlx::{
Describe, Either, Execute, Executor, Postgres, SqlStr,
error::BoxDynError,
postgres::{PgArguments, PgQueryResult, PgRow, PgStatement, PgTypeInfo},
};
// TODO(transact-rs/sqlx#4337): sqlx-postgres 0.9.0 still names every statement when
// statement_cache_capacity is 0, leaving one open per query on the server and colliding
// as "sqlx_s_1 already exists" behind transaction poolers. #4337 makes capacity 0 send
// unnamed statements on its own; delete this wrapper once a sqlx release carries it.
#[derive(Debug)]
pub struct Statements<E> {
executor: E,
named: bool,
}
impl<E> Statements<E> {
pub fn new(executor: E, named: bool) -> Self {
Self { executor, named }
}
}
struct Naming<Q> {
query: Q,
named: bool,
}
impl<'q, Q: Execute<'q, Postgres>> Execute<'q, Postgres> for Naming<Q> {
fn sql(self) -> SqlStr {
self.query.sql()
}
fn statement(&self) -> Option<&PgStatement> {
self.query.statement()
}
fn take_arguments(&mut self) -> Result<Option<PgArguments>, BoxDynError> {
self.query.take_arguments()
}
fn persistent(&self) -> bool {
self.named && self.query.persistent()
}
}
impl<'c, E: Executor<'c, Database = Postgres>> Executor<'c> for Statements<E> {
type Database = Postgres;
fn fetch_many<'e, 'q: 'e, Q>(
self,
query: Q,
) -> BoxStream<'e, Result<Either<PgQueryResult, PgRow>, sqlx::Error>>
where
'c: 'e,
Q: 'q + Execute<'q, Postgres>,
{
self.executor.fetch_many(Naming {
query,
named: self.named,
})
}
fn fetch_optional<'e, 'q: 'e, Q>(
self,
query: Q,
) -> BoxFuture<'e, Result<Option<PgRow>, sqlx::Error>>
where
'c: 'e,
Q: 'q + Execute<'q, Postgres>,
{
self.executor.fetch_optional(Naming {
query,
named: self.named,
})
}
fn prepare_with<'e>(
self,
sql: SqlStr,
parameters: &'e [PgTypeInfo],
) -> BoxFuture<'e, Result<PgStatement, sqlx::Error>>
where
'c: 'e,
{
self.executor.prepare_with(sql, parameters)
}
fn describe<'e>(self, sql: SqlStr) -> BoxFuture<'e, Result<Describe<Postgres>, sqlx::Error>>
where
'c: 'e,
{
self.executor.describe(sql)
}
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use rstest::rstest;
use sqlx::{AssertSqlSafe, PgConnection};
use super::Statements;
use crate::{Endpoint, test_support::database_endpoint};
async fn named_statements(
connection: &mut PgConnection,
named: bool,
) -> Result<i64, sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "pg_prepared_statements is not part of the compile-checked schema"
)]
sqlx::query_scalar("SELECT count(*) FROM pg_prepared_statements WHERE NOT from_sql")
.fetch_one(Statements::new(connection, named))
.await
}
#[rstest]
#[case::cache_off(0, 0)]
#[case::cache_on(100, 4)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn the_endpoint_cache_decides_whether_statements_stay_named(
#[case] statement_cache_capacity: usize,
#[case] expected: i64,
) -> Result<(), Box<dyn StdError>> {
let endpoint = Endpoint {
statement_cache_capacity,
..database_endpoint()?
};
let pool = endpoint.pool()?;
let mut connection = pool.acquire().await?;
for query in ["SELECT 1", "SELECT 2", "SELECT 3"] {
#[expect(clippy::disallowed_methods, reason = "probe queries need no schema")]
sqlx::query(AssertSqlSafe(query))
.execute(Statements::new(
&mut *connection,
endpoint.names_statements(),
))
.await?;
}
assert_eq!(
named_statements(&mut connection, endpoint.names_statements()).await?,
expected
);
Ok(())
}
}

View file

@ -1,10 +1,40 @@
use std::error::Error as StdError;
use std::{error::Error as StdError, str::FromStr};
use litellm_secret_value::SecretValue;
use rstest::fixture;
use sqlx::PgPool;
use sqlx::{PgPool, postgres::PgConnectOptions};
use crate::{Endpoint, PoolLimits, SslMode};
#[fixture]
pub async fn database_pool() -> Result<PgPool, Box<dyn StdError>> {
let database_url = std::env::var("DATABASE_URL")?;
Ok(PgPool::connect(&database_url).await?)
}
pub fn database_endpoint() -> Result<Endpoint, Box<dyn StdError>> {
let database_url = std::env::var("DATABASE_URL")?;
let options = PgConnectOptions::from_str(&database_url)?;
Ok(Endpoint {
host: options.get_host().to_owned(),
port: options.get_port(),
socket: None,
username: options.get_username().to_owned(),
password: url::Url::parse(&database_url)?
.password()
.map(SecretValue::new),
database: options.get_database().unwrap_or("postgres").to_owned(),
ssl_mode: SslMode::Prefer,
ssl_root_cert: None,
ssl_client: None,
application_name: None,
options: Vec::new(),
statement_cache_capacity: 100,
pool: PoolLimits {
max_connections: 4,
acquire_timeout_seconds: Some(10.0),
idle_timeout_seconds: None,
socket_timeout_seconds: None,
},
})
}

View file

@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
rustls.workspace = true
litellm-secret-value.workspace = true
litellm-db.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-spend-clickhouse.workspace = true
@ -71,6 +72,7 @@ serde_json.workspace = true
strum.workspace = true
veil.workspace = true
thiserror.workspace = true
sqlx.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]

View file

@ -1,19 +1,51 @@
use litellm_host_python::ready_future;
use pyo3::{prelude::*, types::PyNone};
use litellm_db::Endpoint;
use litellm_secret_value::SecretValue;
use pyo3::{exceptions::PyValueError, prelude::*};
use sqlx::PgPool;
#[pyclass(frozen)]
pub(crate) struct NativeDatabase {}
pub(crate) struct NativeDatabase {
writer_endpoint: Endpoint,
writer: PgPool,
}
fn map_error(error: litellm_db::Error) -> PyErr {
PyValueError::new_err(litellm_error::Report::new(&error).to_string())
}
// TODO: take typed writer/reader settings and a token provider, then open a sqlx PgPool, once the first
// repository moves to native. The Prisma DATABASE_URL carries Prisma-only params and rotating IAM tokens
#[pymethods]
impl NativeDatabase {
#[staticmethod]
fn connect(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
ready_future(py, Bound::new(py, NativeDatabase {})?.as_any())
fn connect(
py: Python<'_>,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] writer: Endpoint,
) -> PyResult<Bound<'_, PyAny>> {
crate::execution::run_async_value(py, async move {
let pool = writer.pool().map_err(map_error)?;
Ok(NativeDatabase {
writer_endpoint: writer,
writer: pool,
})
})
}
fn set_writer_password(&self, password: String) {
self.writer.set_connect_options(
self.writer_endpoint
.with_password(SecretValue::new(password))
.connect_options(),
);
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
ready_future(py, PyNone::get(py).as_any())
let writer = self.writer.clone();
crate::execution::run_async(
py,
async move {
writer.close().await;
Ok(())
},
|error: PyErr| error,
)
}
}

View file

@ -8,6 +8,7 @@ from typing_extensions import disjoint_base
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
from litellm.redis.stores.router_session_pins.keys import PinKey
from litellm.rust_bridge.database import NativeDatabaseEndpoint
from litellm.rust_bridge.execution import Execution
from litellm.rust_bridge.public_call import NativeCall
@ -562,7 +563,8 @@ class NativeSpendCounters:
@final
class NativeDatabase:
@staticmethod
def connect() -> Future[NativeDatabase]: ...
def connect(writer: NativeDatabaseEndpoint) -> Future[NativeDatabase]: ...
def set_writer_password(self, password: str) -> None: ...
def close(self) -> Future[None]: ...
@final

View file

@ -0,0 +1,33 @@
from typing import Literal, TypeAlias
from typing_extensions import ReadOnly, TypedDict
SslMode: TypeAlias = Literal["disable", "prefer", "require", "verify-ca", "verify-full"]
class NativePoolLimits(TypedDict):
max_connections: ReadOnly[int]
acquire_timeout_seconds: ReadOnly[float | None]
idle_timeout_seconds: ReadOnly[float | None]
socket_timeout_seconds: ReadOnly[float | None]
class NativeClientIdentity(TypedDict):
cert_pem: ReadOnly[str]
key_pem: ReadOnly[str]
class NativeDatabaseEndpoint(TypedDict):
host: ReadOnly[str]
port: ReadOnly[int]
socket: ReadOnly[str | None]
username: ReadOnly[str]
password: ReadOnly[str | None]
database: ReadOnly[str]
ssl_mode: ReadOnly[SslMode]
ssl_root_cert: ReadOnly[str | None]
ssl_client: ReadOnly[NativeClientIdentity | None]
application_name: ReadOnly[str | None]
options: ReadOnly[tuple[tuple[str, str], ...]]
statement_cache_capacity: ReadOnly[int]
pool: ReadOnly[NativePoolLimits]

View file

@ -1,11 +1,15 @@
import asyncio
from typing import Final
from litellm.proxy.db.native_database import native_endpoint
from litellm.rust_bridge._native import NativeDatabase
_NEVER_DIALED_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy"
async def _connect() -> NativeDatabase:
return await NativeDatabase.connect()
async def connect_native_database() -> NativeDatabase:
return await NativeDatabase.connect(native_endpoint(_NEVER_DIALED_URL, tls_offered=True, max_connections=1).endpoint)
def native_database() -> NativeDatabase:
return asyncio.run(_connect())
return asyncio.run(connect_native_database())

View file

@ -0,0 +1,46 @@
from typing import Final
import pytest
from litellm.proxy.db.connection_budget import ConnectionBudget
@pytest.mark.parametrize("total", range(1, 41))
def test_default_split_leaves_prisma_the_whole_total(total: int) -> None:
budget: Final = ConnectionBudget.split(total)
assert (budget.python, budget.rust) == (total, 0)
@pytest.mark.parametrize("total", range(3, 41))
def test_a_reserve_comes_out_of_the_total(total: int) -> None:
budget: Final = ConnectionBudget.split(total, rust_reserved=2)
assert (budget.python + budget.rust, budget.rust) == (total, 2)
@pytest.mark.parametrize(("total", "python"), [(2, 1), (1, 1)], ids=["reserve_equals_total", "single_connection"])
def test_prisma_keeps_one_connection_when_the_reserve_does_not_fit(total: int, python: int) -> None:
budget: Final = ConnectionBudget.split(total, rust_reserved=2)
assert (budget.python, budget.rust) == (python, 2)
@pytest.mark.parametrize("connection_limit", ["7", 7])
def test_a_positive_integer_limit_is_split(connection_limit: str | int) -> None:
assert ConnectionBudget.of(connection_limit, rust_reserved=2) == ConnectionBudget.split(7, rust_reserved=2)
@pytest.mark.parametrize("connection_limit", ["0", 0, "-1", -1, "abc", "10.0", 10.5, "", "\u00b2"])
def test_a_limit_that_is_not_a_positive_integer_is_left_to_prisma(connection_limit: str | float) -> None:
assert ConnectionBudget.of(connection_limit, rust_reserved=2) is None
def test_prisma_params_change_only_beside_a_rust_pool() -> None:
params: Final = {"connection_limit": "010", "pool_timeout": 30}
assert ConnectionBudget.split(10).prisma_params(params) == params
assert ConnectionBudget.split(10, rust_reserved=2).prisma_params(params) == {
"connection_limit": 8,
"pool_timeout": 30,
}

View file

@ -34,6 +34,8 @@ from pydantic import ValidationError
from litellm.proxy.db.db_url_settings import (
PG_SSL_REQUEST,
DatabaseURLSettings,
add_missing_query_params,
prisma_default_params,
token_refresh_params_from_url,
translate_libpq_ssl_params,
unsupported_db_scheme,
@ -1278,32 +1280,52 @@ def test_reader_keeps_its_own_pinned_idle_lifetime(monkeypatch):
@pytest.mark.parametrize(
"writer_limit, reader_limit, num_workers, expected",
"writer_limit, rust_writer_limit, reader_limit, num_workers, expected",
[
("10", None, "4", "4 worker(s) x writer connection_limit 10 = up to 40 connections"),
("10", 0, None, "4", "4 worker(s) x writer connection_limit 10 = up to 40 connections"),
(
"10",
0,
"10",
"4",
"4 worker(s) x (writer connection_limit 10 + reader connection_limit 10) = up to 80 connections",
),
("8", 2, None, "4", "4 worker(s) x (writer connection_limit 8 + rust writer pool 2) = up to 40 connections"),
(
"10",
"8",
2,
"8",
"4",
"4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 8) = up to 72 connections",
),
(
"8",
2,
"50",
"4",
"4 worker(s) x (writer connection_limit 10 + reader connection_limit 50) = up to 240 connections",
"4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 50) = up to 240 connections",
),
("10", None, "0", "1 worker(s) x writer connection_limit 10 = up to 10 connections"),
("8", 2, None, "0", "1 worker(s) x (writer connection_limit 8 + rust writer pool 2) = up to 10 connections"),
],
ids=[
"no_rust_share_writer_only",
"no_rust_share_with_reader",
"writer_only",
"reader_adds_an_engine",
"reader_pins_its_own_limit",
"worker_floor",
],
ids=["writer_only", "reader_doubles_engines", "reader_pins_its_own_limit", "worker_floor"],
)
def test_connection_budget_message_sums_the_engines_each_worker_owns(
writer_limit: str, reader_limit: str | None, num_workers: str, expected: str
writer_limit: str, rust_writer_limit: int, reader_limit: str | None, num_workers: str, expected: str
) -> None:
from litellm.proxy.db.db_url_settings import postgres_connection_budget_message
message: Final = postgres_connection_budget_message(
writer_limit=writer_limit, reader_limit=reader_limit, num_workers=num_workers
writer_limit=writer_limit,
rust_writer_limit=rust_writer_limit,
reader_limit=reader_limit,
num_workers=num_workers,
)
assert expected in message
assert "max_connections minus superuser_reserved_connections" in message
@ -1312,5 +1334,30 @@ def test_connection_budget_message_sums_the_engines_each_worker_owns(
def test_connection_budget_message_survives_an_unparseable_limit() -> None:
from litellm.proxy.db.db_url_settings import postgres_connection_budget_message
message: Final = postgres_connection_budget_message(writer_limit="ten", reader_limit=None, num_workers="2")
message: Final = postgres_connection_budget_message(
writer_limit="ten", rust_writer_limit=2, reader_limit=None, num_workers="2"
)
assert "writer='ten'" in message
@pytest.mark.parametrize(
("url", "native_connections", "expected_name"),
[
("postgresql://u:p@db:5432/proxy", 0, None),
("postgresql://u:p@db:5432/proxy", 2, "litellm-prisma"),
("postgresql://u:p@db:5432/proxy?application_name=ops", 2, "ops"),
],
ids=["untagged_without_native_pool", "tagged_beside_native_pool", "operator_name_wins"],
)
def test_prisma_is_tagged_only_beside_a_native_pool(
url: str, native_connections: int, expected_name: str | None
) -> None:
final_url: Final = add_missing_query_params(
url, prisma_default_params({"max_idle_connection_lifetime": 60}, native_connections)
)
query: Final = urllib.parse.parse_qs(urllib.parse.urlsplit(final_url).query)
assert (query.get("application_name"), query["max_idle_connection_lifetime"]) == (
None if expected_name is None else [expected_name],
["60"],
)

View file

@ -0,0 +1,373 @@
import asyncio
import datetime as dt
from datetime import datetime, timedelta
from pathlib import Path
from typing import Final
import pytest
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.serialization import BestAvailableEncryption, load_pem_private_key, pkcs12
from cryptography.x509.oid import NameOID
from pydantic import TypeAdapter
from litellm.proxy.db.native_database import (
TokenRotation,
keep_writer_token_fresh,
native_endpoint,
needs_tls_probe,
server_offers_tls,
writer_token_rotation,
)
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
from litellm.proxy.db.token_auth import IAMEndpoint, RdsIamTokenAuth
from litellm.rust_bridge.database import NativeDatabaseEndpoint, NativePoolLimits
NOW: Final = datetime(2026, 11, 3, 12, 0, 0)
_ADDRESS: Final = TypeAdapter(tuple[str, int])
PRISMA_DEFAULT_POOL: Final = NativePoolLimits(
max_connections=2, acquire_timeout_seconds=10.0, idle_timeout_seconds=None, socket_timeout_seconds=None
)
def _rds_token(issued: datetime, lifetime_seconds: int) -> str:
return (
"db.internal:5432/?Action=connect&DBUser=litellm"
f"&X-Amz-Date={issued:%Y%m%dT%H%M%SZ}&X-Amz-Expires={lifetime_seconds}&X-Amz-Signature=abc"
)
def _endpoint(**overrides: object) -> NativeDatabaseEndpoint:
return NativeDatabaseEndpoint(
**{ # pyright: ignore[reportArgumentType] # test overrides are checked by the equality assertion
"host": "db",
"port": 5432,
"socket": None,
"username": "u",
"password": None,
"database": "proxy",
"ssl_mode": "prefer",
"ssl_root_cert": None,
"ssl_client": None,
"application_name": "litellm-native",
"options": (),
"statement_cache_capacity": 100,
"pool": PRISMA_DEFAULT_POOL,
**overrides,
}
)
@pytest.mark.parametrize(
"url, tls_offered, expected",
[
pytest.param(
"postgresql://litellm:p%40ss%3Aword@db.internal:6543/proxy?connection_limit=12&pool_timeout=20"
"&max_idle_connection_lifetime=60&socket_timeout=7&application_name=litellm-prisma",
True,
_endpoint(
host="db.internal",
port=6543,
username="litellm",
password="p@ss:word",
pool=NativePoolLimits(
max_connections=2,
acquire_timeout_seconds=20.0,
idle_timeout_seconds=60.0,
socket_timeout_seconds=7.0,
),
),
id="password_pool_params_and_native_share",
),
pytest.param(
"postgresql://svc%40corp@db/proxy?sslmode=require&sslcert=/certs/root.pem&sslaccept=strict"
"&schema=Tenant&options=-c%20statement_timeout%3D30000%20-clock_timeout%3D5000&pgbouncer=true",
True,
_endpoint(
username="svc@corp",
ssl_mode="verify-full",
ssl_root_cert="/certs/root.pem",
options=(("search_path", '"Tenant"'), ("statement_timeout", "30000"), ("lock_timeout", "5000")),
statement_cache_capacity=0,
),
id="strict_tls_schema_timeouts_and_pgbouncer",
),
pytest.param(
"postgresql://u@db/proxy?statement_cache_size=0&application_name=billing-proxy",
True,
_endpoint(statement_cache_capacity=0, application_name="billing-proxy"),
id="statement_cache_size_zero_and_an_operator_application_name_carry_over",
),
pytest.param(
"postgresql://u@db/proxy?statement_cache_size=250&pgbouncer=false",
True,
_endpoint(statement_cache_capacity=250),
id="statement_cache_size_sets_the_capacity",
),
pytest.param(
"postgresql://u@db/proxy?statement_cache_size=250&pgbouncer=true",
True,
_endpoint(statement_cache_capacity=0),
id="pgbouncer_disables_the_statement_cache_whatever_its_size",
),
pytest.param(
"postgresql://u@db/proxy?sslmode=require&sslcert=/certs/root.pem&schema=public",
True,
_endpoint(ssl_mode="require"),
id="require_without_strict_never_verifies_the_operator_cert",
),
pytest.param(
"postgresql://u@db/proxy?sslcert=/certs/root.pem&sslaccept=strict",
True,
_endpoint(ssl_mode="verify-full", ssl_root_cert="/certs/root.pem"),
id="prefer_strict_verifies_when_the_server_offers_tls",
),
pytest.param(
"postgresql://u@db/proxy?sslcert=/certs/root.pem&sslaccept=strict",
False,
_endpoint(ssl_mode="disable"),
id="prefer_strict_falls_back_to_plaintext_without_server_tls",
),
pytest.param(
"postgresql://u@db/proxy?sslmode=disable&sslaccept=strict",
True,
_endpoint(ssl_mode="disable"),
id="disable_wins_over_strict",
),
pytest.param(
"postgresql://u@localhost/proxy?host=/cloudsql/project:region:db&pool_timeout=0",
True,
_endpoint(
host="localhost",
socket="/cloudsql/project:region:db",
pool=NativePoolLimits(
max_connections=2,
acquire_timeout_seconds=None,
idle_timeout_seconds=None,
socket_timeout_seconds=None,
),
),
id="unix_socket_and_pool_timeout_zero_waits_forever",
),
],
)
def test_native_endpoint_reads_the_prisma_url(url: str, tls_offered: bool, expected: NativeDatabaseEndpoint) -> None:
settings: Final = native_endpoint(url, tls_offered=tls_offered, max_connections=2)
assert settings.endpoint == expected
assert settings.ignored_options == ()
@pytest.mark.parametrize(
"options, expected_options, expected_ignored",
[
pytest.param("-c statement_timeout=5 -B 64", (("statement_timeout", "5"),), ("-B", "64"), id="flags"),
pytest.param("statement_timeout=5", (), ("statement_timeout=5",), id="bare_setting"),
pytest.param("-c 'unterminated", (), ("-c 'unterminated",), id="unbalanced_quote"),
],
)
def test_native_endpoint_reports_options_it_cannot_send_instead_of_failing(
options: str, expected_options: tuple[tuple[str, str], ...], expected_ignored: tuple[str, ...]
) -> None:
settings: Final = native_endpoint(f"postgresql://u@db/proxy?options={options}", tls_offered=True, max_connections=2)
assert settings.endpoint["options"] == expected_options
assert settings.ignored_options == expected_ignored
def test_native_endpoint_rejects_a_url_without_a_user() -> None:
with pytest.raises(ValueError, match="DATABASE_URL must name a host and a user"):
native_endpoint("postgresql://db.internal/proxy", tls_offered=True, max_connections=2)
def _identity(password: bytes) -> tuple[bytes, ec.EllipticCurvePrivateKey, x509.Certificate]:
key: Final = ec.generate_private_key(ec.SECP256R1())
name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "litellm")])
cert: Final = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(name)
.public_key(key.public_key())
.serial_number(1)
.not_valid_before(dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc))
.not_valid_after(dt.datetime(2027, 1, 1, tzinfo=dt.timezone.utc))
.sign(key, hashes.SHA256())
)
bundle: Final = pkcs12.serialize_key_and_certificates(
b"litellm", key, cert, None, BestAvailableEncryption(password)
)
return bundle, key, cert
def test_native_endpoint_hands_over_the_pkcs12_identity_as_pem() -> None:
bundle, key, cert = _identity(b"s3cret")
read: Final = {Path("/certs/client.p12"): bundle}
identity: Final = native_endpoint(
"postgresql://u@db/proxy?sslmode=require&sslidentity=/certs/client.p12&sslpassword=s3cret",
tls_offered=True,
max_connections=2,
read_file=read.__getitem__,
).endpoint["ssl_client"]
assert identity is not None
assert x509.load_pem_x509_certificate(identity["cert_pem"].encode()) == cert
assert load_pem_private_key(identity["key_pem"].encode(), None).public_key() == key.public_key()
def test_native_endpoint_rejects_a_wrong_pkcs12_password() -> None:
bundle, _, _ = _identity(b"s3cret")
with pytest.raises(ValueError, match=r"Invalid password or PKCS12 data|Could not deserialize PKCS12 data"):
native_endpoint(
"postgresql://u@db/proxy?sslidentity=/certs/client.p12&sslpassword=wrong",
tls_offered=True,
max_connections=2,
read_file=lambda _path: bundle,
)
@pytest.mark.parametrize(
"url, expected",
[
pytest.param("postgresql://u@db/proxy?sslaccept=strict", True, id="default_prefer_strict"),
pytest.param("postgresql://u@db/proxy?sslmode=prefer&sslaccept=strict", True, id="prefer_strict"),
pytest.param("postgresql://u@db/proxy?sslmode=require&sslaccept=strict", False, id="require_never_falls_back"),
pytest.param("postgresql://u@db/proxy", False, id="prefer_unverified_falls_back_itself"),
pytest.param("postgresql://u@localhost/proxy?sslaccept=strict&host=/run/pg", False, id="unix_socket"),
],
)
def test_only_prefer_strict_needs_a_tls_probe(url: str, expected: bool) -> None:
assert needs_tls_probe(url) is expected
@pytest.mark.asyncio
@pytest.mark.parametrize("answer, offered", [(b"S", True), (b"N", False)])
async def test_server_offers_tls_reads_the_ssl_request_answer(answer: bytes, offered: bool) -> None:
received: Final = asyncio.Event()
async def respond(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
if await reader.readexactly(8) == b"\x00\x00\x00\x08\x04\xd2\x16\x2f":
received.set()
writer.write(answer)
await writer.drain()
writer.close()
server: Final = await asyncio.start_server(respond, "127.0.0.1", 0)
_, port = _ADDRESS.validate_python(server.sockets[0].getsockname())
try:
assert await server_offers_tls(f"postgresql://u@127.0.0.1:{port}/proxy") is offered
finally:
server.close()
await server.wait_closed()
assert received.is_set()
@pytest.mark.asyncio
async def test_unreachable_server_counts_as_offering_tls() -> None:
assert await server_offers_tls("postgresql://u@127.0.0.1:1/proxy", timeout=1) is True
def test_writer_token_rotation_follows_token_auth(monkeypatch: pytest.MonkeyPatch) -> None:
url: Final = "postgresql://litellm:tok@db.internal:5432/proxy?schema=s"
monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False)
monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False)
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "false")
assert writer_token_rotation(url) is None
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv("AWS_RDS_REGION", "us-west-2")
assert writer_token_rotation(url) == TokenRotation(
auth=RdsIamTokenAuth(region="us-west-2"),
endpoint=IAMEndpoint(host="db.internal", port="5432", user="litellm", name="proxy", schema="s"),
)
def test_pooled_writer_needs_no_token_rotation(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "true")
assert writer_token_rotation("postgresql://litellm:local@127.0.0.1:6432/proxy") is None
@pytest.mark.parametrize(
"token, expected_seconds",
[
pytest.param(_rds_token(NOW, 900), 720.0, id="refreshes_buffer_before_expiry"),
pytest.param(_rds_token(NOW - timedelta(seconds=890), 900), 30.0, id="never_sooner_than_retry_floor"),
pytest.param("opaque", 600.0, id="unreadable_expiry_uses_fallback"),
],
)
def test_token_rotation_schedules_the_next_refresh(token: str, expected_seconds: float) -> None:
rotation: Final = TokenRotation(
auth=RdsIamTokenAuth(), endpoint=IAMEndpoint(host="db", port="5432", user="u", name="n")
)
assert rotation.seconds_until_refresh(token, NOW) == expected_seconds
class _Sink:
def __init__(self, wanted: int) -> None:
self.passwords: tuple[str, ...] = ()
self._wanted: Final = wanted
self.done: Final = asyncio.Event()
def set_writer_password(self, password: str) -> None:
self.passwords = (*self.passwords, password)
if len(self.passwords) >= self._wanted:
self.done.set()
class _Schedule:
def __init__(self) -> None:
self.delays: tuple[float, ...] = ()
async def sleep(self, seconds: float) -> None:
self.delays = (*self.delays, seconds)
await asyncio.sleep(0)
def _minting(*outcomes: str | Exception) -> TokenRotation:
remaining: Final = iter(outcomes)
def mint(_auth: object, _endpoint: object) -> str:
outcome: Final = next(remaining)
if isinstance(outcome, Exception):
raise outcome
return outcome
return TokenRotation(
auth=RdsIamTokenAuth(), endpoint=IAMEndpoint(host="db", port="5432", user="u", name="n"), mint_token=mint
)
async def _run_until(sink: _Sink, rotation: TokenRotation, schedule: _Schedule) -> None:
task: Final = asyncio.create_task(
keep_writer_token_fresh(sink, rotation, "opaque", now=lambda: NOW, sleep=schedule.sleep)
)
await asyncio.wait_for(sink.done.wait(), timeout=5)
task.cancel()
await asyncio.gather(task, return_exceptions=True)
@pytest.mark.asyncio
async def test_keep_writer_token_fresh_installs_each_minted_token() -> None:
sink: Final = _Sink(wanted=2)
schedule: Final = _Schedule()
await _run_until(sink, _minting("first%2Ftoken", _rds_token(NOW, 900)), schedule)
assert sink.passwords == ("first/token", _rds_token(NOW, 900))
assert schedule.delays[:2] == (600.0, 600.0)
@pytest.mark.asyncio
async def test_keep_writer_token_fresh_retries_a_failed_mint() -> None:
sink: Final = _Sink(wanted=1)
schedule: Final = _Schedule()
await _run_until(sink, _minting(RuntimeError("sts down"), "recovered"), schedule)
assert sink.passwords == ("recovered",)
assert schedule.delays[:2] == (600.0, 30.0)

View file

@ -32,7 +32,7 @@ from litellm.proxy.proxy_server import (
)
from litellm.repositories.database import Database
from litellm.repositories.deployment_warning_ignore_repository import DEPLOYMENT_WARNING_IGNORES_ROW_ID
from litellm.rust_bridge._native import NativeDatabase
from tests._support.native_database import connect_native_database
T0: Final = datetime(2026, 10, 14, 12, 0, tzinfo=timezone.utc)
THIRTY_DAYS: Final = 30 * 24 * 60 * 60
@ -254,7 +254,7 @@ class TestDeploymentWarningIgnoreStore:
db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh))
)
store: Final = await get_deployment_warning_ignore_repository(
Database(prisma=prisma, native=await NativeDatabase.connect())
Database(prisma=prisma, native=await connect_native_database())
)
assert store is not None
assert await store.get_ignores() == {

View file

@ -16,6 +16,7 @@ import click
import fastapi
import httpx
import pytest
from pydantic import TypeAdapter
import builtins
@ -3518,6 +3519,68 @@ def _run_server_and_capture_urls(
return {k: os.environ[k] for k in _CAPTURED_DB_ENV_VARS if k in os.environ}
_CAPTURED_URLS: Final = TypeAdapter(dict[str, str])
@pytest.mark.parametrize(
"general_settings, database_url, expected_limit, expected_name",
[
pytest.param({}, "postgresql://t:t@localhost:5432/t", "10", None, id="default_limit"),
pytest.param(
{"database_connection_pool_limit": 30},
"postgresql://t:t@localhost:5432/t?application_name=ops",
"30",
"ops",
id="configured_limit_and_operator_name",
),
pytest.param(
{"database_extra_connection_params": {"connection_limit": 5}},
"postgresql://t:t@localhost:5432/t",
"5",
None,
id="extra_params_limit",
),
pytest.param(
{"database_connection_pool_limit": 0},
"postgresql://t:t@localhost:5432/t",
"0",
None,
id="zero_limit_prisma_accepts",
),
pytest.param(
{"database_extra_connection_params": {"connection_limit": "abc"}},
"postgresql://t:t@localhost:5432/t",
"abc",
None,
id="non_integer_limit_left_to_prisma",
),
],
)
def test_prisma_keeps_the_whole_connection_limit_and_its_name(
tmp_path: Path,
general_settings: dict[str, object],
database_url: str,
expected_limit: str,
expected_name: str | None,
) -> None:
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings}))
captured: Final = _CAPTURED_URLS.validate_python(
_run_server_and_capture_urls(
str(config_path), database_url=database_url, direct_url="postgresql://t:t@localhost:5432/t"
)
)
writer: Final = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
direct: Final = urlparse.parse_qs(urlparse.urlparse(captured["DIRECT_URL"]).query)
assert (writer["connection_limit"], writer.get("application_name")) == (
[expected_limit],
None if expected_name is None else [expected_name],
)
assert direct["connection_limit"] == [expected_limit]
class TestReadReplicaConnectionParams:
"""The reader is a second Prisma client with its own pool. Without the
configured params on DATABASE_URL_READ_REPLICA it sizes itself from Prisma's
@ -3552,6 +3615,7 @@ class TestReadReplicaConnectionParams:
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["connection_limit"] == ["3"]
assert "application_name" not in query
assert query["pool_timeout"] == ["20"]
assert query["connect_timeout"] == ["15"]
assert query["socket_timeout"] == ["120"]

View file

@ -1702,6 +1702,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
# Mock the Prisma import
monkeypatch.setattr("litellm.proxy.proxy_server.PrismaClient", MockPrisma)
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
# Create test app
app = FastAPI()
@ -1932,6 +1933,7 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc
_boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
with (
@ -1955,12 +1957,14 @@ async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_
_boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries())
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY)
with pytest.raises(RuntimeError, match="permission denied"):
async with proxy_startup_event(FastAPI()):
pass
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "")
monkeypatch.delenv("DATABASE_URL")
async with proxy_startup_event(FastAPI()) as state:
assert state["database"] is None
@ -2023,6 +2027,7 @@ async def test_proxy_shutdown_stops_database_background_tasks(monkeypatch, tmp_p
yaml.dump({"general_settings": {"master_key": MASTER_KEY}}, f)
monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path))
monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma)
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
monkeypatch.setattr(proxy_server_module, "store_model_in_db", False)
async with proxy_startup_event(FastAPI()):

View file

@ -3,16 +3,80 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
from litellm.repositories.database import Database
WRITER_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy?connection_limit=3"
NATIVE_CONNECTIONS: Final = 2
def _prisma(aclose: AsyncMock | None = None) -> MagicMock:
return MagicMock(aclose=aclose or AsyncMock())
@pytest.fixture
def static_password(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False)
monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False)
monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False)
@pytest.mark.asyncio
@pytest.mark.usefixtures("static_password")
async def test_closing_the_database_closes_prisma() -> None:
prisma: Final = MagicMock()
prisma.aclose = AsyncMock()
database: Final = await Database.open(prisma)
aclose: Final = AsyncMock()
database: Final = await Database.open(_prisma(aclose), WRITER_URL, NATIVE_CONNECTIONS)
prisma.aclose.assert_not_awaited()
aclose.assert_not_awaited()
await database.aclose()
prisma.aclose.assert_awaited_once()
aclose.assert_awaited_once()
def _task_names(database: Database) -> tuple[str, ...]:
return tuple(task.get_coro().__qualname__ for task in database.tasks)
@pytest.mark.asyncio
@pytest.mark.usefixtures("static_password")
async def test_static_password_runs_no_background_task() -> None:
database: Final = await Database.open(_prisma(), WRITER_URL, NATIVE_CONNECTIONS)
assert _task_names(database) == ()
await database.aclose()
@pytest.mark.asyncio
async def test_closing_the_database_stops_the_token_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False)
monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False)
database: Final = await Database.open(_prisma(), WRITER_URL, NATIVE_CONNECTIONS)
assert _task_names(database) == ("keep_writer_token_fresh",)
await database.aclose()
assert all(task.done() for task in database.tasks)
assert database.tasks[0].cancelled()
@pytest.mark.asyncio
async def test_rejected_native_settings_still_close_prisma() -> None:
aclose: Final = AsyncMock()
with pytest.raises(ValueError, match="must name a host and a user"):
await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", NATIVE_CONNECTIONS)
aclose.assert_awaited_once()
@pytest.mark.asyncio
async def test_zero_native_connections_open_no_native_pool(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
aclose: Final = AsyncMock()
database: Final = await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", 0)
assert (database.native, database.tasks) == (None, ())
await database.aclose()
aclose.assert_awaited_once()