diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e7bab3c838c..2a55face8ce 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1136,6 +1136,11 @@ async def fetch_global_spend_with_event_coordination( """ async def _load_global_spend() -> float | None: + from litellm.repositories.native_auth_repository import native_auth_repository + + native: Final = native_auth_repository(prisma_client) + if native is not None: + return await bounded_db_lookup(native.spend("user", LITELLM_PROXY_BUDGET_NAME), name="proxy_budget") proxy_budget_row: Final = await bounded_db_lookup( UserRepository(prisma_client).table.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), name="proxy_budget", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3c2ba2092d0..e807749d7da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -332,7 +332,9 @@ if TYPE_CHECKING: from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction + from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.repositories.prisma_protocols import TableActions + from litellm.rust_bridge._native import NativeDatabase from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep @@ -5675,6 +5677,8 @@ class PrismaClient: http_client: "HttpConfig | None" = None, engine_exit_wait_seconds: float = _SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS, ): + self.native_auth_database: NativeDatabase | None = None + self._native_database_url: Final = database_url ## init logging object self.baseline_accounting_transactions: list[BaselineAccountingRecord] = [] self.baseline_accounting_lock: Final = asyncio.Lock() @@ -5774,6 +5778,7 @@ class PrismaClient: self.db = writer_wrapper else: self.db = writer_wrapper # Client to connect to Prisma db + self._native_reader_url: Final = read_replica_url if isinstance(self.db, RoutingPrismaWrapper) else None self._db_reconnect_lock = asyncio.Lock() self._closing: bool = False self._engine_exit_wait_seconds: Final = engine_exit_wait_seconds @@ -5842,6 +5847,23 @@ class PrismaClient: return self.db.read_target return self.db + @property + def native_writer_settings(self) -> "NativeDatabaseSettings": + from litellm.repositories.database_settings import NativeDatabaseSettings + + return NativeDatabaseSettings.from_url(self._native_database_url, self.writer_db.token_auth) + + def set_native_auth_database(self, native: "NativeDatabase | None") -> None: + self.native_auth_database = native + + @property + def native_reader_settings(self) -> "NativeDatabaseSettings | None": + from litellm.repositories.database_settings import NativeDatabaseSettings + + if self._native_reader_url is None or not isinstance(self.db, RoutingPrismaWrapper): + return None + return NativeDatabaseSettings.from_url(self._native_reader_url, self.db.reader.token_auth) + def tx( self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT, max_wait: timedelta = _PRISMA_DEFAULT_TX_MAX_WAIT ) -> "TransactionManager": diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index d60a2099477..5beb75b1ca6 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -1,9 +1,13 @@ import asyncio from collections.abc import AsyncGenerator, Coroutine from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace +from datetime import datetime, timezone from typing import TYPE_CHECKING, Final +from litellm._logging import verbose_proxy_logger +from litellm.proxy.db.token_auth import parse_database_token_expiration +from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.rust_bridge._native import NativeDatabase if TYPE_CHECKING: @@ -14,23 +18,84 @@ if TYPE_CHECKING: class Database: prisma: "PrismaClient" native: NativeDatabase + password_tasks: tuple[asyncio.Task[None], ...] = () @classmethod async def open(cls, prisma: "PrismaClient") -> "Database": try: - native: Final = await NativeDatabase.connect() + configured: Final = getattr(prisma, "native_writer_settings", None) + writer: Final = ( + await _initial_password(configured) if isinstance(configured, NativeDatabaseSettings) else None + ) + reader_configured: Final = getattr(prisma, "native_reader_settings", None) + reader: Final = ( + await _initial_password(reader_configured) + if isinstance(reader_configured, NativeDatabaseSettings) + else None + ) + native: Final = await NativeDatabase.connect(writer, reader) except BaseException: await prisma.aclose() raise - return cls(prisma=prisma, native=native) + prisma.set_native_auth_database(native if writer is not None else None) + return cls( + prisma=prisma, + native=native, + password_tasks=tuple( + asyncio.create_task(_refresh_password(native, settings, is_reader)) + for settings, is_reader in ((writer, False), (reader, True)) + if settings is not None and settings.password_provider is not None + ), + ) async def aclose(self) -> None: + self.prisma.set_native_auth_database(None) + for task in self.password_tasks: + task.cancel() + await asyncio.gather(*self.password_tasks, return_exceptions=True) try: await self.prisma.aclose() finally: await self.native.close() +async def _initial_password(settings: NativeDatabaseSettings) -> NativeDatabaseSettings: + if settings.password_provider is None: + return settings + return replace(settings, password=await asyncio.to_thread(settings.password_provider)) + + +def _refresh_delay(settings: NativeDatabaseSettings, password: str | None) -> float: + expires: Final = ( + parse_database_token_expiration(settings.token_auth, password) + if settings.token_auth is not None and password is not None + else None + ) + return ( + 600.0 + if expires is None + else max(10.0, (expires - datetime.now(timezone.utc).replace(tzinfo=None)).total_seconds() - 180.0) + ) + + +async def _refresh_password(native: NativeDatabase, settings: NativeDatabaseSettings, reader: bool) -> None: + if settings.password_provider is None: + return + await asyncio.sleep(_refresh_delay(settings, settings.password)) + while True: + try: + password: Final = await asyncio.to_thread(settings.password_provider) + if reader: + native.set_passwords(reader_password=password) + else: + native.set_passwords(writer_password=password) + except Exception: + verbose_proxy_logger.warning("Native database credential refresh failed") + await asyncio.sleep(30.0) + else: + await asyncio.sleep(_refresh_delay(settings, password)) + + @asynccontextmanager async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: task: Final = asyncio.create_task(work) diff --git a/litellm/repositories/database_settings.py b/litellm/repositories/database_settings.py new file mode 100644 index 00000000000..99f9a98ef07 --- /dev/null +++ b/litellm/repositories/database_settings.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass, field +from functools import partial +from typing import Final +from urllib.parse import parse_qsl, unquote, urlsplit + +from litellm.proxy.db.token_auth import DatabaseTokenAuth, mint_database_token, parse_iam_endpoint_from_url + + +class NativeDatabaseSettingsError(ValueError): + def __init__(self, detail: str) -> None: + super().__init__(f"Invalid native database settings. {detail}") + + +def _password(auth: DatabaseTokenAuth, url: str) -> str: + return unquote(mint_database_token(auth, parse_iam_endpoint_from_url(url))) + + +@dataclass(frozen=True, slots=True) +class NativeDatabaseSettings: + host: str + port: int + username: str + password: str | None = field(repr=False) + database: str + schema: str + ssl_mode: str + ssl_root_cert: str | None + max_connections: int + acquire_timeout: float + idle_timeout: float | None + statement_cache_capacity: int + connect_timeout: float + password_provider: Callable[[], str] | None = field(default=None, repr=False) + token_auth: DatabaseTokenAuth | None = field(default=None, repr=False) + + @classmethod + def from_url(cls, url: str, token_auth: DatabaseTokenAuth | None = None) -> NativeDatabaseSettings: + parsed: Final = urlsplit(url) + if parsed.scheme not in ("postgres", "postgresql"): + raise NativeDatabaseSettingsError("PostgreSQL is required") + params: Final = dict(parse_qsl(parsed.query, keep_blank_values=True)) + host: Final = params.get("host", parsed.hostname or "localhost") + ssl_mode: Final = params.get("sslmode", "prefer") + return cls( + host=host, + port=parsed.port or 5432, + username=unquote(parsed.username or "postgres"), + password=None if parsed.password is None else unquote(parsed.password), + database=unquote(parsed.path.lstrip("/") or "postgres"), + schema=params.get("schema", "public"), + ssl_mode="verify-full" if params.get("sslaccept") == "strict" and ssl_mode != "disable" else ssl_mode, + ssl_root_cert=params.get("sslrootcert", params.get("sslcert")), + max_connections=int(params.get("connection_limit", str((os.cpu_count() or 1) * 2 + 1))), + acquire_timeout=float(params.get("pool_timeout", "10")), + idle_timeout=float(params.get("max_idle_connection_lifetime", "300")) or None, + statement_cache_capacity=0 if params.get("pgbouncer", "false").lower() == "true" else 100, + connect_timeout=float(params.get("connect_timeout", "5")), + password_provider=None if token_auth is None else partial(_password, token_auth, url), + token_auth=token_auth, + ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 9af85516ffa..f1b2a2c841d 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -7,6 +7,7 @@ from pydantic import JsonValue from typing_extensions import disjoint_base from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings +from litellm.repositories.database_settings import NativeDatabaseSettings from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall @@ -458,7 +459,15 @@ class NativeSpendCounters: @final class NativeDatabase: @staticmethod - def connect() -> Future[NativeDatabase]: ... + def connect( + writer_settings: NativeDatabaseSettings | None = None, + reader_settings: NativeDatabaseSettings | None = None, + ) -> Future[NativeDatabase]: ... + def auth_read(self, kind: str, ids: list[str], writer: bool = False) -> Future[list[dict[str, JsonValue]]]: ... + def auth_spend(self, kind: str, id: str) -> Future[float | None]: ... + @property + def reader_available(self) -> bool: ... + def set_passwords(self, writer_password: str | None = None, reader_password: str | None = None) -> None: ... def close(self) -> Future[None]: ... @final diff --git a/tests/unit/repositories/test_database_settings.py b/tests/unit/repositories/test_database_settings.py new file mode 100644 index 00000000000..8dac4a6a4cb --- /dev/null +++ b/tests/unit/repositories/test_database_settings.py @@ -0,0 +1,62 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm.repositories.database_settings import NativeDatabaseSettings, NativeDatabaseSettingsError + + +def test_projects_prisma_connection_settings_without_forwarding_url() -> None: + settings: Final = NativeDatabaseSettings.from_url( + "postgresql://auth%40example.test:p%3Aa%2Fs%25s@db.test:5433/gateway%20db" + "?schema=auth%20schema&connection_limit=7&pool_timeout=3&connect_timeout=4" + "&max_idle_connection_lifetime=60&pgbouncer=true&sslmode=require&sslaccept=strict&sslcert=/ca.pem" + ) + assert settings == NativeDatabaseSettings( + host="db.test", + port=5433, + username="auth@example.test", + password="p:a/s%s", + database="gateway db", + schema="auth schema", + ssl_mode="verify-full", + ssl_root_cert="/ca.pem", + max_connections=7, + acquire_timeout=3, + idle_timeout=60, + statement_cache_capacity=0, + connect_timeout=4, + ) + + +@pytest.mark.parametrize( + ("query", "ssl_mode", "certificate"), + ( + ("sslmode=disable&sslaccept=strict", "disable", None), + ("sslmode=verify-ca&sslrootcert=/root.pem", "verify-ca", "/root.pem"), + ("sslmode=require&sslaccept=accept_invalid_certs", "require", None), + ("sslmode=prefer", "prefer", None), + ), +) +def test_preserves_tls_policy(query: str, ssl_mode: str, certificate: str | None) -> None: + settings: Final = NativeDatabaseSettings.from_url(f"postgresql://user:password@db.test/gateway?{query}") + assert (settings.ssl_mode, settings.ssl_root_cert) == (ssl_mode, certificate) + + +def test_database_settings_repr_redacts_password_and_provider() -> None: + settings: Final = NativeDatabaseSettings.from_url("postgresql://user:opaque-password@db.test/gateway") + assert "opaque-password" not in repr(settings) + assert repr(replace(settings, password="another-password")) == repr(settings) + + +def test_preserves_unix_socket_and_disabled_idle_expiration() -> None: + settings: Final = NativeDatabaseSettings.from_url( + "postgresql://user@localhost/gateway?host=%2Fvar%2Frun%2Fpostgresql&max_idle_connection_lifetime=0" + ) + assert (settings.host, settings.password, settings.idle_timeout) == ("/var/run/postgresql", None, None) + + +def test_rejects_non_postgres_without_echoing_credentials() -> None: + with pytest.raises(NativeDatabaseSettingsError, match="PostgreSQL is required") as error: + NativeDatabaseSettings.from_url("mysql://user:opaque-password@db.test/gateway") + assert "opaque-password" not in str(error.value)