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()