From dc097107b2d24f94fb1f819be75e2338ed9f7a55 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 13 Nov 2026 11:04:45 +0900 Subject: [PATCH] feat(db): add the SQLx pool foundation with the writer pool 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 --- .github/workflows/test-rust-db.yml | 2 + litellm/proxy/db/connection_budget.py | 52 +++ litellm/proxy/db/db_url_settings.py | 36 +- litellm/proxy/db/native_database.py | 281 +++++++++++++ litellm/proxy/db/pgbouncer.py | 28 +- litellm/proxy/proxy_cli.py | 21 +- litellm/proxy/proxy_server.py | 5 +- litellm/repositories/database.py | 42 +- litellm/rust/Cargo.lock | 4 + litellm/rust/crates/db/AGENTS.md | 35 ++ litellm/rust/crates/db/Cargo.toml | 7 +- litellm/rust/crates/db/queries/.gitattributes | 1 + litellm/rust/crates/db/src/deadline.rs | 51 +++ litellm/rust/crates/db/src/endpoint.rs | 332 ++++++++++++++++ litellm/rust/crates/db/src/error.rs | 2 + litellm/rust/crates/db/src/lib.rs | 6 + litellm/rust/crates/db/src/statements.rs | 152 +++++++ litellm/rust/crates/db/src/test_support.rs | 34 +- litellm/rust/crates/python-bridge/Cargo.toml | 2 + litellm/rust/crates/python-bridge/src/db.rs | 48 ++- litellm/rust_bridge/_native.pyi | 4 +- litellm/rust_bridge/database.py | 33 ++ tests/_support/native_database.py | 10 +- tests/unit/proxy/db/test_connection_budget.py | 46 +++ tests/unit/proxy/db/test_db_url_settings.py | 65 ++- tests/unit/proxy/db/test_native_database.py | 373 ++++++++++++++++++ .../test_deployment_warning_endpoints.py | 4 +- tests/unit/proxy/test_proxy_cli.py | 64 +++ ...test_proxy_server_endpoints_and_startup.py | 5 + tests/unit/repositories/test_database.py | 74 +++- 30 files changed, 1757 insertions(+), 62 deletions(-) create mode 100644 litellm/proxy/db/connection_budget.py create mode 100644 litellm/proxy/db/native_database.py create mode 100644 litellm/rust/crates/db/queries/.gitattributes create mode 100644 litellm/rust/crates/db/src/deadline.rs create mode 100644 litellm/rust/crates/db/src/endpoint.rs create mode 100644 litellm/rust/crates/db/src/statements.rs create mode 100644 litellm/rust_bridge/database.py create mode 100644 tests/unit/proxy/db/test_connection_budget.py create mode 100644 tests/unit/proxy/db/test_native_database.py diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml index 116d6785287..3a8af7f170f 100644 --- a/.github/workflows/test-rust-db.yml +++ b/.github/workflows/test-rust-db.yml @@ -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 diff --git a/litellm/proxy/db/connection_budget.py b/litellm/proxy/db/connection_budget.py new file mode 100644 index 00000000000..71b839b91ef --- /dev/null +++ b/litellm/proxy/db/connection_budget.py @@ -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}) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 9db060d5c9d..231904d4d78 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -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" diff --git a/litellm/proxy/db/native_database.py b/litellm/proxy/db/native_database.py new file mode 100644 index 00000000000..eb5f5ff40b8 --- /dev/null +++ b/litellm/proxy/db/native_database.py @@ -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) diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py index c9fbabd3585..d9983ce774a 100644 --- a/litellm/proxy/db/pgbouncer.py +++ b/litellm/proxy/db/pgbouncer.py @@ -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( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index f252e69d739..cef24e441c1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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}" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f9d4d5704d8..65853f61549 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index d60a2099477..b577ba2bd67 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -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 diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index ea790ecfb61..ddd0eaf0870 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -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", diff --git a/litellm/rust/crates/db/AGENTS.md b/litellm/rust/crates/db/AGENTS.md index 1e08317230e..1fedf9c9a2b 100644 --- a/litellm/rust/crates/db/AGENTS.md +++ b/litellm/rust/crates/db/AGENTS.md @@ -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 | diff --git a/litellm/rust/crates/db/Cargo.toml b/litellm/rust/crates/db/Cargo.toml index 4f8fe4e28e2..527a0927338 100644 --- a/litellm/rust/crates/db/Cargo.toml +++ b/litellm/rust/crates/db/Cargo.toml @@ -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 diff --git a/litellm/rust/crates/db/queries/.gitattributes b/litellm/rust/crates/db/queries/.gitattributes new file mode 100644 index 00000000000..cc2d335b830 --- /dev/null +++ b/litellm/rust/crates/db/queries/.gitattributes @@ -0,0 +1 @@ +*.sql text eol=lf diff --git a/litellm/rust/crates/db/src/deadline.rs b/litellm/rust/crates/db/src/deadline.rs new file mode 100644 index 00000000000..d986daa108c --- /dev/null +++ b/litellm/rust/crates/db/src/deadline.rs @@ -0,0 +1,51 @@ +use std::{io, time::Duration}; + +use crate::Error; + +pub async fn within( + socket_timeout: Option, + operation: impl Future>, +) -> Result { + 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)); + } +} diff --git a/litellm/rust/crates/db/src/endpoint.rs b/litellm/rust/crates/db/src/endpoint.rs new file mode 100644 index 00000000000..cb04943333b --- /dev/null +++ b/litellm/rust/crates/db/src/endpoint.rs @@ -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, + pub username: String, + pub password: Option, + pub database: String, + pub ssl_mode: SslMode, + pub ssl_root_cert: Option, + pub ssl_client: Option, + pub application_name: Option, + 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, + pub idle_timeout_seconds: Option, + pub socket_timeout_seconds: Option, +} + +impl From 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 { + 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, Error> { + seconds(self.pool.socket_timeout_seconds) + } +} + +fn seconds(value: Option) -> Result, 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> { + 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> { + 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> { + 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(()) + } +} diff --git a/litellm/rust/crates/db/src/error.rs b/litellm/rust/crates/db/src/error.rs index af0ee030896..c4d44a3231e 100644 --- a/litellm/rust/crates/db/src/error.rs +++ b/litellm/rust/crates/db/src/error.rs @@ -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 { diff --git a/litellm/rust/crates/db/src/lib.rs b/litellm/rust/crates/db/src/lib.rs index 541f97c7db9..bab2d492c92 100644 --- a/litellm/rust/crates/db/src/lib.rs +++ b/litellm/rust/crates/db/src/lib.rs @@ -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; diff --git a/litellm/rust/crates/db/src/statements.rs b/litellm/rust/crates/db/src/statements.rs new file mode 100644 index 00000000000..a3da065fe81 --- /dev/null +++ b/litellm/rust/crates/db/src/statements.rs @@ -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 { + executor: E, + named: bool, +} + +impl Statements { + pub fn new(executor: E, named: bool) -> Self { + Self { executor, named } + } +} + +struct Naming { + query: Q, + named: bool, +} + +impl<'q, Q: Execute<'q, Postgres>> Execute<'q, Postgres> for Naming { + fn sql(self) -> SqlStr { + self.query.sql() + } + + fn statement(&self) -> Option<&PgStatement> { + self.query.statement() + } + + fn take_arguments(&mut self) -> Result, 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 { + type Database = Postgres; + + fn fetch_many<'e, 'q: 'e, Q>( + self, + query: Q, + ) -> BoxStream<'e, Result, 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, 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> + where + 'c: 'e, + { + self.executor.prepare_with(sql, parameters) + } + + fn describe<'e>(self, sql: SqlStr) -> BoxFuture<'e, Result, 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 { + #[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> { + 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(()) + } +} diff --git a/litellm/rust/crates/db/src/test_support.rs b/litellm/rust/crates/db/src/test_support.rs index 992be4a4447..40446ca53dc 100644 --- a/litellm/rust/crates/db/src/test_support.rs +++ b/litellm/rust/crates/db/src/test_support.rs @@ -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> { let database_url = std::env::var("DATABASE_URL")?; Ok(PgPool::connect(&database_url).await?) } + +pub fn database_endpoint() -> Result> { + 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, + }, + }) +} diff --git a/litellm/rust/crates/python-bridge/Cargo.toml b/litellm/rust/crates/python-bridge/Cargo.toml index 48206d242e4..504706def5f 100644 --- a/litellm/rust/crates/python-bridge/Cargo.toml +++ b/litellm/rust/crates/python-bridge/Cargo.toml @@ -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] diff --git a/litellm/rust/crates/python-bridge/src/db.rs b/litellm/rust/crates/python-bridge/src/db.rs index 03fb30b986f..ca1ccc8fb2b 100644 --- a/litellm/rust/crates/python-bridge/src/db.rs +++ b/litellm/rust/crates/python-bridge/src/db.rs @@ -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> { - 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> { + 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> { - 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, + ) } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 111b771e65f..964862db17a 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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 diff --git a/litellm/rust_bridge/database.py b/litellm/rust_bridge/database.py new file mode 100644 index 00000000000..1888a349f68 --- /dev/null +++ b/litellm/rust_bridge/database.py @@ -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] diff --git a/tests/_support/native_database.py b/tests/_support/native_database.py index 29b1fff399f..bade92d79f8 100644 --- a/tests/_support/native_database.py +++ b/tests/_support/native_database.py @@ -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()) diff --git a/tests/unit/proxy/db/test_connection_budget.py b/tests/unit/proxy/db/test_connection_budget.py new file mode 100644 index 00000000000..115bba5d68b --- /dev/null +++ b/tests/unit/proxy/db/test_connection_budget.py @@ -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, + } diff --git a/tests/unit/proxy/db/test_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index 1f86f6ab5d1..a42894efa89 100644 --- a/tests/unit/proxy/db/test_db_url_settings.py +++ b/tests/unit/proxy/db/test_db_url_settings.py @@ -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"], + ) diff --git a/tests/unit/proxy/db/test_native_database.py b/tests/unit/proxy/db/test_native_database.py new file mode 100644 index 00000000000..dad278ea13c --- /dev/null +++ b/tests/unit/proxy/db/test_native_database.py @@ -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) diff --git a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py index 67aea8dfeb8..456cc2005e9 100644 --- a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py @@ -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() == { diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index d85d685c352..4e2ddfe859d 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -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"] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 09a896ec053..051c578b2e9 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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()): diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index d9343a796d2..798bdc39325 100644 --- a/tests/unit/repositories/test_database.py +++ b/tests/unit/repositories/test_database.py @@ -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()