feat(auth): own native database pools in proxy lifecycle

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",

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":

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)

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,
+        )

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

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)