From fb4d32efea6503ad3de0e6a8a167d38be5ff57dd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:05:36 +0000 Subject: [PATCH] feat(proxy): add Azure passwordless datastore auth Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/_redis.py | 127 +++++-------- litellm/_redis_credential_provider.py | 6 +- litellm/proxy/auth/azure_postgres_token.py | 22 +++ litellm/proxy/db/db_url_settings.py | 82 ++++++--- litellm/proxy/db/prisma_client.py | 168 ++++++++++++------ litellm/proxy/db/routing_prisma_wrapper.py | 2 +- litellm/proxy/proxy_cli.py | 38 ++-- litellm/proxy/utils.py | 34 ++-- .../get_azure_ad_token_provider.py | 54 +++++- .../proxy/auth/test_azure_postgres_token.py | 17 ++ .../proxy/db/test_db_url_settings.py | 85 ++++----- .../proxy/db/test_rds_iam_token_expiry.py | 50 +++++- .../proxy/db/test_routing_prisma_wrapper.py | 109 ++++++------ .../test_get_azure_ad_token_provider.py | 32 +++- tests/test_litellm/test_redis.py | 152 ++++++++-------- 15 files changed, 592 insertions(+), 386 deletions(-) create mode 100644 litellm/proxy/auth/azure_postgres_token.py create mode 100644 tests/test_litellm/proxy/auth/test_azure_postgres_token.py diff --git a/litellm/_redis.py b/litellm/_redis.py index 0b91cdabffc..bd7c2e3589a 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -29,10 +29,15 @@ from litellm.constants import ( REDIS_SOCKET_TIMEOUT, ) from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.secret_managers.get_azure_ad_token_provider import ( + AzureTokenCredential, + build_azure_identity_credential, +) from ._logging import verbose_logger AZURE_REDIS_SCOPE = "https://redis.azure.com/.default" +_REDIS_CREDENTIAL_PROVIDER_KEY = "_litellm_credential_provider" def _get_redis_kwargs(): @@ -64,7 +69,7 @@ def _get_redis_kwargs(): def _get_redis_url_kwargs(client=None): if client is None: client = redis.Redis.from_url - arg_spec = inspect.getfullargspec(redis.Redis.from_url) + arg_spec = inspect.getfullargspec(client) # Only allow primitive arguments exclude_args = { @@ -73,7 +78,7 @@ def _get_redis_url_kwargs(client=None): "retry", } - include_args = ["url"] + include_args = ["url", "credential_provider"] available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args @@ -178,38 +183,18 @@ def _build_azure_credential( azure_client_id: Optional[str] = None, azure_tenant_id: Optional[str] = None, azure_client_secret: Optional[str] = None, -): +) -> AzureTokenCredential: """ Build a long-lived Azure credential object. Azure SDK credentials cache tokens internally and handle expiry/refresh transparently, so this should be called once and the result reused. """ - try: - from azure.identity import ( - ClientSecretCredential, - DefaultAzureCredential, - ManagedIdentityCredential, - ) - except ImportError: - raise ImportError( - "azure-identity is required for Azure AD Redis authentication. Install it with: pip install azure-identity" - ) - - _client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID") - _tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID") - _client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET") - - if _client_id and _tenant_id and _client_secret: - return ClientSecretCredential( - client_id=_client_id, - tenant_id=_tenant_id, - client_secret=_client_secret, - ) - elif _client_id: - return ManagedIdentityCredential(client_id=_client_id) - else: - return DefaultAzureCredential() + return build_azure_identity_credential( + azure_client_id=azure_client_id, + azure_tenant_id=azure_tenant_id, + azure_client_secret=azure_client_secret, + ) def _generate_azure_ad_redis_token( @@ -238,6 +223,7 @@ def create_azure_ad_redis_connect_func( azure_client_id: Optional[str] = None, azure_tenant_id: Optional[str] = None, azure_client_secret: Optional[str] = None, + credential: AzureTokenCredential | None = None, ) -> Callable: """ Creates a custom Redis connection function for Azure AD authentication. @@ -246,7 +232,7 @@ def create_azure_ad_redis_connect_func( closure) and reused across connections — the Azure SDK handles token caching and silent renewal internally. Only ``get_token`` is called per connection. """ - credential = _build_azure_credential( + azure_credential = credential or _build_azure_credential( azure_client_id=azure_client_id, azure_tenant_id=azure_tenant_id, azure_client_secret=azure_client_secret, @@ -262,7 +248,7 @@ def create_azure_ad_redis_connect_func( self._parser.on_connect(self) - access_token = credential.get_token(AZURE_REDIS_SCOPE).token + access_token = azure_credential.get_token(AZURE_REDIS_SCOPE).token # Only include username when explicitly set — sending AUTH "" # is invalid for most ACL-configured Azure Redis instances. @@ -284,11 +270,6 @@ def create_azure_ad_redis_connect_func( if str_if_bytes(auth_response) != "OK": raise AuthenticationError("Azure AD authentication failed for Redis") - # Attach the live credential object so async paths can wrap it in - # AzureADCredentialProvider for refresh-aware token retrieval. The raw - # client_id/tenant_id/secret are intentionally NOT exposed here — the - # credential closure already holds them. - ad_connect._azure_credential = credential # type: ignore[attr-defined] return ad_connect @@ -382,8 +363,7 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func( service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs ) - # Store GCP service account in redis_connect_func for async cluster access - redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account # type: ignore[attr-defined] + redis_kwargs[_REDIS_CREDENTIAL_PROVIDER_KEY] = GCPIAMCredentialProvider(_gcp_service_account) # Remove GCP-specific kwargs that shouldn't be passed to Redis client redis_kwargs.pop("gcp_service_account", None) @@ -410,17 +390,17 @@ def _get_redis_client_logic(**env_overrides): _azure_client_secret = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET") verbose_logger.debug("Setting up Azure AD authentication for Redis.") - redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func( + azure_credential = _build_azure_credential( azure_client_id=_azure_client_id, azure_tenant_id=_azure_tenant_id, azure_client_secret=_azure_client_secret, ) - # Marker for async paths to detect Azure AD auth. The live credential - # object is attached separately as `_azure_credential` by - # `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret - # are intentionally NOT exposed on the function to avoid leaking - # credentials via inspection or logging. - redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True # type: ignore[attr-defined] + redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(credential=azure_credential) + username = redis_kwargs.get("username") + redis_kwargs[_REDIS_CREDENTIAL_PROVIDER_KEY] = AzureADCredentialProvider( + azure_credential, + username=str(username) if username is not None else None, + ) # Always remove Azure-specific kwargs that shouldn't be passed to Redis client redis_kwargs.pop("azure_redis_ad_token", None) @@ -538,11 +518,15 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: def get_redis_client(**env_overrides): redis_kwargs = _get_redis_client_logic(**env_overrides) + credential_provider = redis_kwargs.pop(_REDIS_CREDENTIAL_PROVIDER_KEY, None) if "startup_nodes" in redis_kwargs: return init_redis_cluster(redis_kwargs) if "url" in redis_kwargs and redis_kwargs["url"] is not None: + redis_kwargs.pop("redis_connect_func", None) + if credential_provider is not None: + redis_kwargs["credential_provider"] = credential_provider args = _get_redis_url_kwargs() url_kwargs = {} for arg in redis_kwargs: @@ -563,6 +547,7 @@ def get_redis_async_client( **env_overrides, ) -> Union[async_redis.Redis, async_redis.RedisCluster]: redis_kwargs = _get_redis_client_logic(**env_overrides) + credential_provider = redis_kwargs.pop(_REDIS_CREDENTIAL_PROVIDER_KEY, None) if "startup_nodes" in redis_kwargs: from redis.cluster import ClusterNode @@ -573,22 +558,9 @@ def get_redis_async_client( if arg in args: cluster_kwargs[arg] = redis_kwargs[arg] - # Handle GCP IAM authentication for async clusters - redis_connect_func = cluster_kwargs.pop("redis_connect_func", None) - - # Use a CredentialProvider so the IAM token is regenerated on every new - # connection — mirrors the sync path where redis_connect_func is invoked - # per connection. Without this, the token would expire after ~1 hour. - if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) - # Handle Azure AD authentication for async clusters via CredentialProvider - # so the credential's internal cache + silent refresh runs per connection - # (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry). - elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"): - cluster_kwargs["credential_provider"] = AzureADCredentialProvider( - redis_connect_func._azure_credential, - username=os.environ.get("REDIS_USERNAME") or None, - ) + cluster_kwargs.pop("redis_connect_func", None) + if credential_provider is not None: + cluster_kwargs["credential_provider"] = credential_provider new_startup_nodes: List[ClusterNode] = [] @@ -614,6 +586,9 @@ def get_redis_async_client( if "url" in redis_kwargs and redis_kwargs["url"] is not None: if connection_pool is not None: return async_redis.Redis(connection_pool=connection_pool) + redis_kwargs.pop("redis_connect_func", None) + if credential_provider is not None: + redis_kwargs["credential_provider"] = credential_provider args = _get_redis_url_kwargs(client=async_redis.Redis.from_url) url_kwargs = {} for arg in redis_kwargs: @@ -629,18 +604,9 @@ def get_redis_async_client( if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs: return _init_async_redis_sentinel(redis_kwargs) - # Wrap GCP / Azure AD auth in a CredentialProvider for the standard async - # Redis client. The async client doesn't support redis_connect_func, but it - # does honour credential_provider — which is called per connection, so the - # underlying SDK can refresh tokens silently before they expire. - redis_connect_func = redis_kwargs.pop("redis_connect_func", None) - if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"): - redis_kwargs["credential_provider"] = AzureADCredentialProvider( - redis_connect_func._azure_credential, - username=os.environ.get("REDIS_USERNAME") or None, - ) - elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) + redis_kwargs.pop("redis_connect_func", None) + if credential_provider is not None: + redis_kwargs["credential_provider"] = credential_provider _pretty_print_redis_config(redis_kwargs=redis_kwargs) @@ -656,6 +622,7 @@ def get_redis_connection_pool( **env_overrides, ) -> Optional[async_redis.BlockingConnectionPool]: redis_kwargs = _get_redis_client_logic(**env_overrides) + credential_provider = redis_kwargs.pop(_REDIS_CREDENTIAL_PROVIDER_KEY, None) verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs) if "startup_nodes" in redis_kwargs: @@ -666,6 +633,8 @@ def get_redis_connection_pool( "timeout": REDIS_CONNECTION_POOL_TIMEOUT, "url": redis_kwargs["url"], } + if credential_provider is not None: + pool_kwargs["credential_provider"] = credential_provider if "max_connections" in redis_kwargs: try: pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"]) @@ -676,17 +645,9 @@ def get_redis_connection_pool( ) return async_redis.BlockingConnectionPool.from_url(**pool_kwargs) - # Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed - # connections re-fetch tokens via the SDK's internal cache + silent refresh - # rather than reusing a single token captured at pool creation. - redis_connect_func = redis_kwargs.pop("redis_connect_func", None) - if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"): - redis_kwargs["credential_provider"] = AzureADCredentialProvider( - redis_connect_func._azure_credential, - username=os.environ.get("REDIS_USERNAME") or None, - ) - elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"): - redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account) + redis_kwargs.pop("redis_connect_func", None) + if credential_provider is not None: + redis_kwargs["credential_provider"] = credential_provider connection_class = async_redis.Connection if redis_kwargs.pop("ssl", False): diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index b973e292a17..15fd673c145 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,10 +1,12 @@ import asyncio import threading import time -from typing import Any, Dict, Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union from redis.credentials import CredentialProvider # type: ignore[attr-defined] +from litellm.secret_managers.get_azure_ad_token_provider import AzureTokenCredential + # Azure AD scope for Redis Cache for Azure. AZURE_REDIS_SCOPE = "https://redis.azure.com/.default" @@ -115,7 +117,7 @@ class AzureADCredentialProvider(CredentialProvider): fail authentication after the initial token expired (~1 hour TTL). """ - def __init__(self, credential: Any, username: Optional[str] = None) -> None: + def __init__(self, credential: AzureTokenCredential, username: Optional[str] = None) -> None: self._credential = credential self._username = username diff --git a/litellm/proxy/auth/azure_postgres_token.py b/litellm/proxy/auth/azure_postgres_token.py new file mode 100644 index 00000000000..cf943d4ee71 --- /dev/null +++ b/litellm/proxy/auth/azure_postgres_token.py @@ -0,0 +1,22 @@ +from urllib.parse import quote + +from litellm.secret_managers.get_azure_ad_token_provider import ( + AzureTokenCredential, + build_azure_identity_credential, +) + +AZURE_POSTGRES_SCOPE = "https://ossrdbms-aad.database.windows.net/.default" + + +def generate_azure_postgres_auth_token( + credential: AzureTokenCredential | None = None, + azure_client_id: str | None = None, + azure_tenant_id: str | None = None, + azure_client_secret: str | None = None, +) -> str: + azure_credential = credential or build_azure_identity_credential( + azure_client_id=azure_client_id, + azure_tenant_id=azure_tenant_id, + azure_client_secret=azure_client_secret, + ) + return quote(azure_credential.get_token(AZURE_POSTGRES_SCOPE).token, safe="") diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 4f220916a4a..606ae5c54a9 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -37,11 +37,10 @@ from typing import Final, cast from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict -# Imported as a module (not `from ... import generate_iam_auth_token`) so the -# AWS-touching token mint stays patchable at its canonical location in tests. -from litellm.proxy.auth import rds_iam_token +from litellm.proxy.auth import azure_postgres_token, rds_iam_token _IAM_ENV_KEY = "IAM_TOKEN_DB_AUTH" +_AZURE_POSTGRES_ENV_KEY = "AZURE_POSTGRESQL_AUTH" _DEFAULT_PG_PORT = "5432" # schema.prisma pins `provider = "postgresql"`, so these are the only schemes @@ -91,6 +90,7 @@ class DatabaseURLSettings(BaseSettings): model_config = SettingsConfigDict(case_sensitive=False, extra="ignore") iam_token_db_auth: bool = Field(default=False, validation_alias=_IAM_ENV_KEY) + azure_postgresql_auth: bool = Field(default=False, validation_alias=_AZURE_POSTGRES_ENV_KEY) # Writer database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") @@ -129,7 +129,8 @@ class DatabaseURLSettings(BaseSettings): enabled but a required field is missing — the proxy cannot recover from this and a clear startup error beats a Prisma connect failure. """ - if self.iam_token_db_auth: + self._validate_token_auth() + if self.iam_token_db_auth or self.azure_postgresql_auth: missing = [ env for env, val in ( @@ -141,20 +142,30 @@ class DatabaseURLSettings(BaseSettings): ] if missing: raise RuntimeError( - "IAM_TOKEN_DB_AUTH is enabled but required DB env var(s) " + "Database token auth is enabled but required DB env var(s) " f"are unset: {', '.join(missing)}. Set them so the writer " - "DATABASE_URL can be assembled with a minted IAM token." + "DATABASE_URL can be assembled with a minted database token." ) host = cast(str, self.database_host) user = cast(str, self.database_user) name = cast(str, self.database_name) - # IAM token is already URL-quoted by generate_iam_auth_token; - # user/name embedded raw (parity with proxy_cli.py / IAMEndpoint). - token = rds_iam_token.generate_iam_auth_token(db_host=host, db_port=self.database_port, db_user=user) - url = f"postgresql://{user}:{token}@{host}:{self.database_port}/{name}" - if self.database_schema: - url += f"?schema={self.database_schema}" - return url + token = ( + azure_postgres_token.generate_azure_postgres_auth_token() + if self.azure_postgresql_auth + else rds_iam_token.generate_iam_auth_token( + db_host=host, + db_port=self.database_port, + db_user=user, + ) + ) + return self._token_url( + user=user, + token=token, + host=host, + port=self.database_port, + name=name, + schema=self.database_schema, + ) # Password auth: an operator-pinned DATABASE_URL always wins. if self.database_url: @@ -177,6 +188,7 @@ class DatabaseURLSettings(BaseSettings): pre-existing ``DATABASE_URL_READ_REPLICA``. Reader fields fall back to the writer's values. """ + self._validate_token_auth() if not self.database_host_read_replica: return None # reader is opt-in if self.database_url_read_replica: @@ -189,7 +201,7 @@ class DatabaseURLSettings(BaseSettings): schema = self.database_schema_read_replica or self.database_schema password = self.database_password_read_replica or self.database_password - if self.iam_token_db_auth: + if self.iam_token_db_auth or self.azure_postgresql_auth: missing = [ env for env, val in ( @@ -200,7 +212,7 @@ class DatabaseURLSettings(BaseSettings): ] if missing: raise RuntimeError( - "IAM_TOKEN_DB_AUTH is enabled and DATABASE_HOST_READ_REPLICA " + "Database token auth is enabled and DATABASE_HOST_READ_REPLICA " "is set, but the reader could not resolve: " f"{', '.join(missing)} (no *_READ_REPLICA value and no " "writer fallback). Set the reader fields or the writer " @@ -208,11 +220,19 @@ class DatabaseURLSettings(BaseSettings): ) user = cast(str, user) name = cast(str, name) - token = rds_iam_token.generate_iam_auth_token(db_host=host, db_port=port, db_user=user) - url = f"postgresql://{user}:{token}@{host}:{port}/{name}" - if schema: - url += f"?schema={schema}" - return url + token = ( + azure_postgres_token.generate_azure_postgres_auth_token() + if self.azure_postgresql_auth + else rds_iam_token.generate_iam_auth_token(db_host=host, db_port=port, db_user=user) + ) + return self._token_url( + user=user, + token=token, + host=host, + port=port, + name=name, + schema=schema, + ) if user and name: return self._password_url( @@ -225,6 +245,26 @@ class DatabaseURLSettings(BaseSettings): ) return None + @staticmethod + def _token_url( + *, + user: str, + token: str, + host: str, + port: str, + name: str, + schema: str | None, + ) -> str: + quote = urllib.parse.quote + url = f"postgresql://{quote(user, safe='')}:{token}@{host}:{port}/{quote(name, safe='')}" + if schema: + url += f"?schema={quote(schema, safe='')}" + return url + + def _validate_token_auth(self) -> None: + if self.iam_token_db_auth and self.azure_postgresql_auth: + raise RuntimeError("IAM_TOKEN_DB_AUTH and AZURE_POSTGRESQL_AUTH cannot both be enabled") + @staticmethod def _password_url( *, @@ -287,6 +327,8 @@ class DatabaseURLSettings(BaseSettings): # Normalize the toggle so downstream readers (PrismaWrapper's # IAM refresh) reliably see IAM on, regardless of spelling. os.environ[_IAM_ENV_KEY] = "True" + if self.azure_postgresql_auth: + os.environ[_AZURE_POSTGRES_ENV_KEY] = "True" wrote_writer = True reader_url = self.build_reader_url() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index fbccb8a726c..4d67f348c12 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -3,6 +3,8 @@ This file contains the PrismaWrapper class, which is used to wrap the Prisma cli """ import asyncio +import base64 +import binascii import os import random import signal @@ -11,13 +13,39 @@ import time import urllib import urllib.parse from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone +from enum import Enum from typing import Any, Callable, Union +from pydantic import BaseModel, ValidationError + from litellm._logging import verbose_proxy_logger from litellm.secret_managers.main import str_to_bool +class DatabaseTokenAuth(str, Enum): + RDS_IAM = "rds_iam" + AZURE_ENTRA = "azure_entra" + + +class AzureTokenClaims(BaseModel): + exp: int + + +def resolve_database_token_auth( + *, + iam_token_db_auth: bool, + azure_postgresql_auth: bool, +) -> DatabaseTokenAuth | None: + if iam_token_db_auth and azure_postgresql_auth: + raise ValueError("IAM_TOKEN_DB_AUTH and AZURE_POSTGRESQL_AUTH cannot both be enabled") + if azure_postgresql_auth: + return DatabaseTokenAuth.AZURE_ENTRA + if iam_token_db_auth: + return DatabaseTokenAuth.RDS_IAM + return None + + @dataclass(frozen=True) class IAMEndpoint: """Static parts of an RDS IAM-authenticated Postgres connection. @@ -34,9 +62,11 @@ class IAMEndpoint: schema: str | None = None def build_url(self, token: str) -> str: - url = f"postgresql://{self.user}:{token}@{self.host}:{self.port}/{self.name}" + user = urllib.parse.quote(self.user, safe="") + name = urllib.parse.quote(self.name, safe="") + url = f"postgresql://{user}:{token}@{self.host}:{self.port}/{name}" if self.schema: - url += f"?schema={self.schema}" + url += f"?schema={urllib.parse.quote(self.schema, safe='')}" return url @@ -62,12 +92,31 @@ def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint: return IAMEndpoint( host=parsed.hostname, port=port, - user=parsed.username, - name=name, - schema=schema, + user=urllib.parse.unquote(parsed.username), + name=urllib.parse.unquote(name), + schema=urllib.parse.unquote(schema) if schema is not None else None, ) +def build_database_token_auth_url( + endpoint: IAMEndpoint, + database_token_auth: DatabaseTokenAuth, +) -> str: + if database_token_auth == DatabaseTokenAuth.AZURE_ENTRA: + from litellm.proxy.auth.azure_postgres_token import generate_azure_postgres_auth_token + + token = generate_azure_postgres_auth_token() + else: + from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token + + token = generate_iam_auth_token( + db_host=endpoint.host, + db_port=endpoint.port, + db_user=endpoint.user, + ) + return endpoint.build_url(token) + + class PrismaWrapper: """ Wrapper around Prisma client that handles RDS IAM token authentication. @@ -97,9 +146,11 @@ class PrismaWrapper: iam_endpoint: IAMEndpoint | None = None, recreate_uses_datasource: bool = False, log_prefix: str = "", + database_token_auth: DatabaseTokenAuth | None = None, ): self._original_prisma = original_prisma - self.iam_token_db_auth = iam_token_db_auth + self.database_token_auth = database_token_auth or (DatabaseTokenAuth.RDS_IAM if iam_token_db_auth else None) + self.iam_token_db_auth = self.database_token_auth is not None # Per-connection knobs so the same wrapper can be used for the writer # (defaults: DATABASE_URL env, IAM endpoint from DATABASE_HOST/etc., @@ -136,6 +187,11 @@ class PrismaWrapper: self._engine_generation: int = 0 self.on_engine_replaced: Callable[[], None] | None = None + def _token_auth_log_name(self) -> str: + if self.database_token_auth == DatabaseTokenAuth.AZURE_ENTRA: + return "Azure PostgreSQL Entra token" + return "RDS IAM token" + def _get_engine_pid(self) -> int: """Get the PID of the current Prisma engine subprocess, or 0 if unavailable. @@ -222,26 +278,28 @@ class PrismaWrapper: if token is None: return None - try: - # Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&... - if "?" not in token: - return None - + if "?" in token: query_string = token.split("?", 1)[1] params = urllib.parse.parse_qs(query_string) - expires_str = params.get("X-Amz-Expires", [None])[0] date_str = params.get("X-Amz-Date", [None])[0] + if expires_str and date_str: + try: + token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ") + return token_created + timedelta(seconds=int(expires_str)) + except ValueError as exc: + verbose_proxy_logger.debug("Failed to parse RDS IAM token expiration: %s", exc) - if not expires_str or not date_str: - return None - - token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ") - expires_in = int(expires_str) - - return token_created + timedelta(seconds=expires_in) - except Exception as e: - verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}") + parts = token.split(".") + if len(parts) < 2: + return None + padding = "=" * (-len(parts[1]) % 4) + try: + payload = base64.urlsafe_b64decode(f"{parts[1]}{padding}") + claims = AzureTokenClaims.model_validate_json(payload) + return datetime.fromtimestamp(claims.exp, tz=timezone.utc).replace(tzinfo=None) + except (binascii.Error, UnicodeDecodeError, ValidationError, ValueError) as exc: + verbose_proxy_logger.debug("Failed to parse Azure PostgreSQL token expiration: %s", exc) return None def _calculate_seconds_until_refresh(self) -> float: @@ -304,30 +362,22 @@ class PrismaWrapper: if not self.iam_token_db_auth: return None - from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token - if self._iam_endpoint is not None: endpoint = self._iam_endpoint - token = generate_iam_auth_token(db_host=endpoint.host, db_port=endpoint.port, db_user=endpoint.user) - _db_url = endpoint.build_url(token) else: - db_host = os.getenv("DATABASE_HOST") - # Default to the Postgres standard port; passing None to - # `generate_iam_auth_token` makes botocore embed the literal - # string "None" in the presigned URL, which then fails to parse. - db_port = os.getenv("DATABASE_PORT", "5432") - db_user = os.getenv("DATABASE_USER") - db_name = os.getenv("DATABASE_NAME") - db_schema = os.getenv("DATABASE_SCHEMA") + endpoint = IAMEndpoint( + host=os.getenv("DATABASE_HOST") or "", + port=os.getenv("DATABASE_PORT", "5432"), + user=os.getenv("DATABASE_USER") or "", + name=os.getenv("DATABASE_NAME") or "", + schema=os.getenv("DATABASE_SCHEMA"), + ) - token = generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user) - - _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" - if db_schema: - _db_url += f"?schema={db_schema}" - - os.environ[self._db_url_env_var] = _db_url - return _db_url + if self.database_token_auth is None: + return None + db_url = build_database_token_auth_url(endpoint, self.database_token_auth) + os.environ[self._db_url_env_var] = db_url + return db_url async def recreate_prisma_client( self, @@ -443,8 +493,9 @@ class PrismaWrapper: self._token_refresh_task = asyncio.create_task(self._token_refresh_loop()) verbose_proxy_logger.info( - "%sStarted RDS IAM token proactive refresh background task", + "%sStarted %s proactive refresh background task", self._log_prefix, + self._token_auth_log_name(), ) async def stop_token_refresh_task(self) -> None: @@ -462,7 +513,7 @@ class PrismaWrapper: except asyncio.CancelledError: pass self._token_refresh_task = None - verbose_proxy_logger.info("%sStopped RDS IAM token refresh background task", self._log_prefix) + verbose_proxy_logger.info("%sStopped %s refresh background task", self._log_prefix, self._token_auth_log_name()) async def _token_refresh_loop(self) -> None: """ @@ -473,7 +524,7 @@ class PrismaWrapper: This is more efficient than polling, requiring only 1 wake-up per token cycle. """ verbose_proxy_logger.info( - f"{self._log_prefix}RDS IAM token refresh loop started. " + f"{self._log_prefix}{self._token_auth_log_name()} refresh loop started. " f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration." ) @@ -484,21 +535,25 @@ class PrismaWrapper: if sleep_seconds > 0: verbose_proxy_logger.info( - f"{self._log_prefix}RDS IAM token refresh scheduled in " + f"{self._log_prefix}{self._token_auth_log_name()} refresh scheduled in " f"{sleep_seconds:.0f} seconds ({sleep_seconds / 60:.1f} minutes)" ) await asyncio.sleep(sleep_seconds) # Refresh the token - verbose_proxy_logger.info("%sProactively refreshing RDS IAM token...", self._log_prefix) + verbose_proxy_logger.info( + "%sProactively refreshing %s...", + self._log_prefix, + self._token_auth_log_name(), + ) await self._safe_refresh_token() except asyncio.CancelledError: - verbose_proxy_logger.info("%sRDS IAM token refresh loop cancelled", self._log_prefix) + verbose_proxy_logger.info("%s%s refresh loop cancelled", self._log_prefix, self._token_auth_log_name()) break except Exception as e: verbose_proxy_logger.error( - f"{self._log_prefix}Error in RDS IAM token refresh loop: {e}. " + f"{self._log_prefix}Error in {self._token_auth_log_name()} refresh loop: {e}. " f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..." ) # On error, wait before retrying to avoid tight error loops @@ -522,8 +577,9 @@ class PrismaWrapper: # by skipping when the current token still has comfortable runway. if self._token_refresh_not_needed(os.getenv(self._db_url_env_var)): verbose_proxy_logger.debug( - "%sRDS IAM token still fresh; skipping redundant refresh.", + "%s%s still fresh; skipping redundant refresh.", self._log_prefix, + self._token_auth_log_name(), ) return @@ -534,13 +590,15 @@ class PrismaWrapper: await self._recreate_prisma_client_locked(new_db_url) self._last_refresh_time = datetime.utcnow() verbose_proxy_logger.info( - "%sRDS IAM token refreshed successfully. New token valid for ~15 minutes.", + "%s%s refreshed successfully.", self._log_prefix, + self._token_auth_log_name(), ) else: verbose_proxy_logger.error( - "%sFailed to generate new RDS IAM token during proactive refresh", + "%sFailed to generate new %s during proactive refresh", self._log_prefix, + self._token_auth_log_name(), ) def _token_refresh_not_needed(self, token_url: str | None) -> bool: @@ -594,10 +652,11 @@ class PrismaWrapper: if running_loop is not None: verbose_proxy_logger.warning( - "%sRDS IAM token expired in __getattr__ — proactive refresh " + "%s%s expired in __getattr__ - proactive refresh " "may have failed. Scheduling async refresh; the current " "request may fail and be retried with the fresh token.", self._log_prefix, + self._token_auth_log_name(), ) # Non-blocking: schedule the locked refresh on the # running loop. The reconnection lock inside @@ -605,9 +664,10 @@ class PrismaWrapper: running_loop.create_task(self._safe_refresh_token()) else: verbose_proxy_logger.warning( - "%sRDS IAM token expired in __getattr__ — proactive refresh " + "%s%s expired in __getattr__ - proactive refresh " "may have failed. Triggering synchronous fallback refresh...", self._log_prefix, + self._token_auth_log_name(), ) new_db_url = self.get_rds_iam_token() if new_db_url: diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 703abc64d8c..c4ac124e865 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -243,7 +243,7 @@ class RoutingPrismaWrapper: if self._reader.iam_token_db_auth: new_reader_url = self._reader.get_rds_iam_token() if not new_reader_url: - raise RuntimeError("Failed to generate fresh IAM token for read replica") + raise RuntimeError("Failed to generate fresh database auth token for read replica") await self._reader.recreate_prisma_client(new_reader_url, http_client=http_client) return reader_url = os.getenv("DATABASE_URL_READ_REPLICA", "") diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 74ec0cc8700..6995e5c1e48 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -703,6 +703,12 @@ class ProxyInitializationHelpers: is_flag=True, help="Connects to RDS DB with IAM token", ) +@click.option( + "--azure_postgresql_auth", + default=False, + is_flag=True, + help="Connects to Azure PostgreSQL with Microsoft Entra token auth", +) @click.option( "--num_requests", default=10, @@ -850,6 +856,7 @@ def run_server( granian_threads, test_async, iam_token_db_auth, + azure_postgresql_auth, num_requests, use_queue, health, @@ -978,29 +985,16 @@ def run_server( general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### - if iam_token_db_auth or get_secret_bool("IAM_TOKEN_DB_AUTH"): - from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token + iam_auth_enabled = iam_token_db_auth or get_secret_bool("IAM_TOKEN_DB_AUTH") + azure_auth_enabled = azure_postgresql_auth or get_secret_bool("AZURE_POSTGRESQL_AUTH") + if iam_auth_enabled or azure_auth_enabled: + from litellm.proxy.db.db_url_settings import DatabaseURLSettings - db_host = os.getenv("DATABASE_HOST") - # Default to the Postgres standard port. Without a default, - # `db_port=None` flows into `boto.generate_db_auth_token(Port=None)` - # and botocore stringifies it to `"None"` while building the - # presigned URL, which then blows up with `ValueError: Port could - # not be cast to integer value as 'None'` during signing. - db_port = os.getenv("DATABASE_PORT", "5432") - db_user = os.getenv("DATABASE_USER") - db_name = os.getenv("DATABASE_NAME") - db_schema = os.getenv("DATABASE_SCHEMA") - - token = generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user) - - # print(f"token: {token}") - _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" - if db_schema: - _db_url += f"?schema={db_schema}" - - os.environ["DATABASE_URL"] = _db_url - os.environ["IAM_TOKEN_DB_AUTH"] = "True" + if iam_auth_enabled: + os.environ["IAM_TOKEN_DB_AUTH"] = "True" + if azure_auth_enabled: + os.environ["AZURE_POSTGRESQL_AUTH"] = "True" + DatabaseURLSettings.from_env().apply_to_env() ### DECRYPT ENV VAR ### diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9f36e729330..52fb53b40fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -128,7 +128,9 @@ from litellm.proxy.db.exception_handler import ( from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import ( PrismaWrapper, + build_database_token_auth_url, parse_iam_endpoint_from_url, + resolve_database_token_auth, ) from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( @@ -2790,6 +2792,7 @@ class PrismaClient: ## init logging object self.proxy_logging_obj = proxy_logging_obj self.iam_token_db_auth: Optional[bool] = str_to_bool(os.getenv("IAM_TOKEN_DB_AUTH")) + azure_postgresql_auth = str_to_bool(os.getenv("AZURE_POSTGRESQL_AUTH")) is True verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma # type: ignore @@ -2799,6 +2802,11 @@ class PrismaClient: verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.") raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.") iam_flag = self.iam_token_db_auth if self.iam_token_db_auth is not None else False + database_token_auth = resolve_database_token_auth( + iam_token_db_auth=iam_flag, + azure_postgresql_auth=azure_postgresql_auth, + ) + token_auth_enabled = database_token_auth is not None # When read-replica routing is on, tag log lines with [writer]/[reader] # so the two wrappers' interleaved IAM refresh logs can be told apart. # Single-DB deployments get an empty prefix (logs unchanged). @@ -2807,14 +2815,16 @@ class PrismaClient: if http_client is not None: writer_wrapper = PrismaWrapper( original_prisma=Prisma(http=http_client), - iam_token_db_auth=iam_flag, + iam_token_db_auth=token_auth_enabled, log_prefix=writer_log_prefix, + database_token_auth=database_token_auth, ) else: writer_wrapper = PrismaWrapper( original_prisma=Prisma(), - iam_token_db_auth=iam_flag, + iam_token_db_auth=token_auth_enabled, log_prefix=writer_log_prefix, + database_token_auth=database_token_auth, ) # Optional read-replica routing. When DATABASE_URL_READ_REPLICA is set, @@ -2829,7 +2839,7 @@ class PrismaClient: # the same cadence as the writer. We parse the static endpoint # pieces (host/port/user/db) once from the reader URL — only # the IAM token rotates after that. - reader_iam_endpoint = parse_iam_endpoint_from_url(read_replica_url) if iam_flag else None + reader_iam_endpoint = parse_iam_endpoint_from_url(read_replica_url) if token_auth_enabled else None # Mint a fresh IAM token for the reader BEFORE constructing the # Prisma client. Mirrors what `proxy_cli.py` already does for # the writer (proxy_cli.py:812-832) — without this, the reader @@ -2838,17 +2848,8 @@ class PrismaClient: # to the synchronous fallback path in # `PrismaWrapper.__getattr__`, which deadlocks the event loop # and times out after 30s. - if iam_flag and reader_iam_endpoint is not None: - from litellm.proxy.auth.rds_iam_token import ( - generate_iam_auth_token, - ) - - reader_token = generate_iam_auth_token( - db_host=reader_iam_endpoint.host, - db_port=reader_iam_endpoint.port, - db_user=reader_iam_endpoint.user, - ) - read_replica_url = reader_iam_endpoint.build_url(reader_token) + if database_token_auth is not None and reader_iam_endpoint is not None: + read_replica_url = build_database_token_auth_url(reader_iam_endpoint, database_token_auth) os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url reader_kwargs: Dict[str, Any] = {"datasource": {"url": read_replica_url}} if http_client is not None: @@ -2857,16 +2858,17 @@ class PrismaClient: reader_prisma = Prisma(**reader_kwargs) reader_wrapper = PrismaWrapper( original_prisma=reader_prisma, - iam_token_db_auth=iam_flag, + iam_token_db_auth=token_auth_enabled, db_url_env_var="DATABASE_URL_READ_REPLICA", iam_endpoint=reader_iam_endpoint, recreate_uses_datasource=True, log_prefix="[reader]", + database_token_auth=database_token_auth, ) self.db = RoutingPrismaWrapper(writer=writer_wrapper, reader=reader_wrapper) verbose_proxy_logger.info( "PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA" - + (" (with IAM token auto-refresh)" if iam_flag else "") + + (" (with database token auto-refresh)" if token_auth_enabled else "") ) except Exception as e: # Reader is opt-in; never let its construction fail proxy diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 2c52054964d..c5221ddda0c 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,5 +1,5 @@ import os -from typing import Any, Callable, Optional, Union +from typing import Any, Callable, Optional, Protocol, Union from litellm._logging import verbose_logger from litellm.types.secret_managers.get_azure_ad_token_provider import ( @@ -7,6 +7,58 @@ from litellm.types.secret_managers.get_azure_ad_token_provider import ( ) +class AzureAccessToken(Protocol): + @property + def token(self) -> str: ... + + @property + def expires_on(self) -> int: ... + + +class AzureTokenCredential(Protocol): + def get_token(self, *scopes: str) -> AzureAccessToken: ... + + +def build_azure_identity_credential( + azure_client_id: str | None = None, + azure_tenant_id: str | None = None, + azure_client_secret: str | None = None, +) -> AzureTokenCredential: + try: + from azure.identity import ( + ClientSecretCredential, + DefaultAzureCredential, + ManagedIdentityCredential, + WorkloadIdentityCredential, + ) + except ImportError as exc: + raise ImportError( + "azure-identity is required for Azure passwordless authentication. " + "Install it with: pip install azure-identity" + ) from exc + + client_id = azure_client_id or os.getenv("AZURE_CLIENT_ID") + tenant_id = azure_tenant_id or os.getenv("AZURE_TENANT_ID") + client_secret = azure_client_secret or os.getenv("AZURE_CLIENT_SECRET") + federated_token_file = os.getenv("AZURE_FEDERATED_TOKEN_FILE") + + if client_id and tenant_id and client_secret: + return ClientSecretCredential( + client_id=client_id, + tenant_id=tenant_id, + client_secret=client_secret, + ) + if client_id and tenant_id and federated_token_file: + return WorkloadIdentityCredential( + client_id=client_id, + tenant_id=tenant_id, + token_file_path=federated_token_file, + ) + if client_id: + return ManagedIdentityCredential(client_id=client_id) + return DefaultAzureCredential() + + def infer_credential_type_from_environment() -> AzureCredentialType: if ( os.environ.get("AZURE_CLIENT_ID") diff --git a/tests/test_litellm/proxy/auth/test_azure_postgres_token.py b/tests/test_litellm/proxy/auth/test_azure_postgres_token.py new file mode 100644 index 00000000000..cd9b90c5ce3 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_azure_postgres_token.py @@ -0,0 +1,17 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from litellm.proxy.auth.azure_postgres_token import ( + AZURE_POSTGRES_SCOPE, + generate_azure_postgres_auth_token, +) + + +def test_generate_azure_postgres_auth_token_uses_database_scope_and_url_encodes_token(): + credential = MagicMock() + credential.get_token.return_value = SimpleNamespace(token="header.payload/signature+padding=", expires_on=0) + + token = generate_azure_postgres_auth_token(credential=credential) + + assert token == "header.payload%2Fsignature%2Bpadding%3D" + credential.get_token.assert_called_once_with(AZURE_POSTGRES_SCOPE) diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index e5aa09addab..e6da2d4e312 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -30,6 +30,7 @@ def _apply() -> bool: _MANAGED_DB_ENV_VARS = ( "IAM_TOKEN_DB_AUTH", + "AZURE_POSTGRESQL_AUTH", "DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA", @@ -96,14 +97,42 @@ def test_assembles_writer_url_when_iam_enabled(monkeypatch): with _stub_iam_token("WRITER_TOKEN"): assert _apply() is True - assert ( - os.environ["DATABASE_URL"] - == "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://litellm:WRITER_TOKEN@writer.example.com:5432/litellm_db" # Reader was never configured, so it must not have been set. assert "DATABASE_URL_READ_REPLICA" not in os.environ +def test_assembles_azure_writer_and_reader_urls(monkeypatch): + monkeypatch.setenv("AZURE_POSTGRESQL_AUTH", "true") + monkeypatch.setenv("DATABASE_HOST", "writer.postgres.database.azure.com") + monkeypatch.setenv("DATABASE_USER", "user@example.com") + monkeypatch.setenv("DATABASE_NAME", "litellm db") + monkeypatch.setenv("DATABASE_HOST_READ_REPLICA", "reader.postgres.database.azure.com") + + with patch( + "litellm.proxy.auth.azure_postgres_token.generate_azure_postgres_auth_token", + side_effect=("WRITER_TOKEN", "READER_TOKEN"), + ): + assert _apply() is True + + assert ( + os.environ["DATABASE_URL"] + == "postgresql://user%40example.com:WRITER_TOKEN@writer.postgres.database.azure.com:5432/litellm%20db" + ) + assert ( + os.environ["DATABASE_URL_READ_REPLICA"] + == "postgresql://user%40example.com:READER_TOKEN@reader.postgres.database.azure.com:5432/litellm%20db" + ) + + +def test_rejects_multiple_database_token_auth_modes(monkeypatch): + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv("AZURE_POSTGRESQL_AUTH", "true") + + with pytest.raises(RuntimeError, match="cannot both be enabled"): + _apply() + + def test_missing_writer_envs_raises(monkeypatch): monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") # DATABASE_HOST intentionally unset. @@ -147,10 +176,7 @@ def test_reader_url_not_clobbered_when_already_set(monkeypatch): with _stub_iam_token("READER_TOKEN"): _apply() - assert ( - os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://app:secret@reader.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://app:secret@reader.example.com:5432/litellm_db" def test_reader_url_skipped_when_host_unset(monkeypatch): @@ -196,10 +222,7 @@ def test_assembles_writer_url_from_password(monkeypatch): monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") assert _apply() is True - assert ( - os.environ["DATABASE_URL"] - == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" def test_writer_password_is_percent_encoded(monkeypatch): @@ -209,28 +232,20 @@ def test_writer_password_is_percent_encoded(monkeypatch): monkeypatch.setenv("DATABASE_PASSWORD", "p@ss/w:rd") assert _apply() is True - assert ( - os.environ["DATABASE_URL"] - == "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://litellm:p%40ss%2Fw%3Ard@writer.example.com:5432/litellm_db" def test_writer_url_not_clobbered_when_already_set(monkeypatch): """An operator-pinned DATABASE_URL (e.g. helm's $(VAR) assembly) always wins over the discrete fields.""" - monkeypatch.setenv( - "DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db" - ) + monkeypatch.setenv("DATABASE_URL", "postgresql://pinned:url@db.example.com:5432/litellm_db") monkeypatch.setenv("DATABASE_HOST", "writer.example.com") monkeypatch.setenv("DATABASE_USER", "litellm") monkeypatch.setenv("DATABASE_NAME", "litellm_db") monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") assert _apply() is False - assert ( - os.environ["DATABASE_URL"] - == "postgresql://pinned:url@db.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://pinned:url@db.example.com:5432/litellm_db" def test_writer_url_passwordless(monkeypatch): @@ -239,10 +254,7 @@ def test_writer_url_passwordless(monkeypatch): monkeypatch.setenv("DATABASE_NAME", "litellm_db") assert _apply() is True - assert ( - os.environ["DATABASE_URL"] - == "postgresql://litellm@writer.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://litellm@writer.example.com:5432/litellm_db" def test_database_username_alias(monkeypatch): @@ -254,10 +266,7 @@ def test_database_username_alias(monkeypatch): monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t") assert _apply() is True - assert ( - os.environ["DATABASE_URL"] - == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL"] == "postgresql://litellm:s3cr3t@writer.example.com:5432/litellm_db" def test_password_reader_falls_back_to_writer_password(monkeypatch): @@ -268,10 +277,7 @@ def test_password_reader_falls_back_to_writer_password(monkeypatch): monkeypatch.setenv("DATABASE_HOST_READ_REPLICA", "reader.example.com") assert _apply() is True - assert ( - os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db" def test_password_reader_uses_own_credentials(monkeypatch): @@ -284,10 +290,7 @@ def test_password_reader_uses_own_credentials(monkeypatch): monkeypatch.setenv("DATABASE_PASSWORD_READ_REPLICA", "ro_pw") assert _apply() is True - assert ( - os.environ["DATABASE_URL_READ_REPLICA"] - == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" - ) + assert os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" @pytest.mark.parametrize( @@ -352,9 +355,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") - monkeypatch.setenv( - "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" - ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db") with pytest.raises(RuntimeError, match="DATABASE_URL_READ_REPLICA.*mysql"): _apply() diff --git a/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py index ca24f856022..1b81d8490d8 100644 --- a/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py +++ b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py @@ -14,6 +14,8 @@ Run these tests: """ import asyncio +import base64 +import json import os import urllib.parse from datetime import datetime, timedelta @@ -45,9 +47,7 @@ class TestPrismaWrapperTokenRefresh: def _set_database_url_with_token(self, expires_in_seconds: int = 900): """Set DATABASE_URL with a mock token.""" token = self._generate_mock_token(expires_in_seconds) - os.environ["DATABASE_URL"] = ( - f"postgresql://test_user:{token}@test-host:5432/test_db" - ) + os.environ["DATABASE_URL"] = f"postgresql://test_user:{token}@test-host:5432/test_db" @pytest.mark.asyncio async def test_is_token_expired_fresh(self, setup_env): @@ -73,9 +73,7 @@ class TestPrismaWrapperTokenRefresh: # Create an expired token old_date = datetime.utcnow() - timedelta(seconds=901) date_str = old_date.strftime("%Y%m%dT%H%M%SZ") - token = ( - f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc" - ) + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc" encoded_token = urllib.parse.quote(token, safe="") db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db" @@ -155,6 +153,42 @@ class TestTokenExpirationParsing: assert wrapper._parse_token_expiration("no-query-params") is None assert wrapper._parse_token_expiration("?missing=params") is None + def test_parse_azure_postgres_jwt_expiration(self): + from litellm.proxy.db.prisma_client import PrismaWrapper + + expires_on = 1893456000 + payload = base64.urlsafe_b64encode(json.dumps({"exp": expires_on}).encode()).decode().rstrip("=") + token = f"header.{payload}.signature" + wrapper = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=True) + + assert wrapper._parse_token_expiration(token) == datetime.utcfromtimestamp(expires_on) + + def test_azure_postgres_refresh_builds_database_url(self, unset_database_url): + from litellm.proxy.db.prisma_client import DatabaseTokenAuth, IAMEndpoint, PrismaWrapper + + wrapper = PrismaWrapper( + original_prisma=MagicMock(), + iam_token_db_auth=True, + iam_endpoint=IAMEndpoint( + host="server.postgres.database.azure.com", + port="5432", + user="user@example.com", + name="litellm db", + ), + database_token_auth=DatabaseTokenAuth.AZURE_ENTRA, + ) + + with patch( + "litellm.proxy.auth.azure_postgres_token.generate_azure_postgres_auth_token", + return_value="AZURE_TOKEN", + ): + database_url = wrapper.get_rds_iam_token() + + assert ( + database_url + == "postgresql://user%40example.com:AZURE_TOKEN@server.postgres.database.azure.com:5432/litellm%20db" + ) + class TestBackgroundRefreshLoop: """Tests for the background refresh loop timing.""" @@ -210,9 +244,7 @@ async def demonstrate_fix(): date_str = now.strftime("%Y%m%dT%H%M%SZ") token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123" encoded_token = urllib.parse.quote(token, safe="") - os.environ["DATABASE_URL"] = ( - f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm" - ) + os.environ["DATABASE_URL"] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm" # Create mock prisma client mock_prisma = MagicMock() diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index e5bb8b99507..8f0fbefb3dd 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -295,12 +295,8 @@ async def test_recreate_prisma_client_recreates_both_writer_and_reader(): with patch.dict(os.environ, {"DATABASE_URL_READ_REPLICA": "reader-url"}): await routing.recreate_prisma_client("writer-url", http_client=None) - writer.recreate_prisma_client.assert_awaited_once_with( - "writer-url", http_client=None, expected_generation=None - ) - reader.recreate_prisma_client.assert_awaited_once_with( - "reader-url", http_client=None - ) + writer.recreate_prisma_client.assert_awaited_once_with("writer-url", http_client=None, expected_generation=None) + reader.recreate_prisma_client.assert_awaited_once_with("reader-url", http_client=None) assert routing.reader_unavailable is False @@ -335,9 +331,7 @@ async def test_recreate_degrades_reader_if_reader_recreate_fails(): writer.recreate_prisma_client = AsyncMock() reader = MagicMock() reader.iam_token_db_auth = False - reader.recreate_prisma_client = AsyncMock( - side_effect=RuntimeError("reader still down") - ) + reader.recreate_prisma_client = AsyncMock(side_effect=RuntimeError("reader still down")) routing = RoutingPrismaWrapper(writer=writer, reader=reader) @@ -390,9 +384,7 @@ async def test_recreate_iam_reader_refreshes_token(): await routing.recreate_prisma_client("writer-url") reader.get_rds_iam_token.assert_called_once() - reader.recreate_prisma_client.assert_awaited_once_with( - "postgresql://u:fresh@h:5432/db", http_client=None - ) + reader.recreate_prisma_client.assert_awaited_once_with("postgresql://u:fresh@h:5432/db", http_client=None) assert routing.reader_unavailable is False @@ -479,9 +471,7 @@ async def test_writer_recreate_passes_http_client_through(monkeypatch): writer = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=False) sentinel_http = object() - await writer.recreate_prisma_client( - "postgresql://u:p@h:5432/db", http_client=sentinel_http - ) + await writer.recreate_prisma_client("postgresql://u:p@h:5432/db", http_client=sentinel_http) assert captured_kwargs == {"http": sentinel_http} @@ -559,14 +549,8 @@ async def test_iam_refresh_logs_carry_log_prefix(caplog): messages = [r.getMessage() for r in caplog.records] # Both start and stop notifications carry the prefix. - assert any( - m.startswith("[reader] Started RDS IAM token proactive refresh") - for m in messages - ) - assert any( - m.startswith("[reader] Stopped RDS IAM token refresh background task") - for m in messages - ) + assert any(m.startswith("[reader] Started RDS IAM token proactive refresh") for m in messages) + assert any(m.startswith("[reader] Stopped RDS IAM token refresh background task") for m in messages) def test_get_rds_iam_token_returns_none_when_iam_disabled(): @@ -697,9 +681,7 @@ def test_writer_get_rds_iam_token_uses_database_host_env_vars(monkeypatch, unset "port": "5432", "user": "litellm", } - assert new_url == ( - "postgresql://litellm:WRITER-TOKEN@writer.aurora.local:5432/litellm?schema=public" - ) + assert new_url == ("postgresql://litellm:WRITER-TOKEN@writer.aurora.local:5432/litellm?schema=public") # Writer updates its own env var (DATABASE_URL by default), not the reader's. assert os.environ["DATABASE_URL"] == new_url @@ -750,9 +732,7 @@ def test_reader_iam_refresh_uses_parsed_endpoint(monkeypatch): "user": "lit", } assert new_url is not None - assert new_url.startswith( - "postgresql://lit:FRESH-TOKEN@reader.aurora.local:5432/litellm" - ) + assert new_url.startswith("postgresql://lit:FRESH-TOKEN@reader.aurora.local:5432/litellm") # The reader updates its OWN env var; writer's DATABASE_URL is left alone. assert os.environ["DATABASE_URL_READ_REPLICA"] == new_url assert os.environ["DATABASE_URL"] == "writer-url-untouched" @@ -786,13 +766,9 @@ async def test_reader_recreate_uses_datasource_override(monkeypatch): recreate_uses_datasource=True, ) - await reader.recreate_prisma_client( - "postgresql://u:newtoken@h:5432/db", http_client=None - ) + await reader.recreate_prisma_client("postgresql://u:newtoken@h:5432/db", http_client=None) - assert captured_kwargs == { - "datasource": {"url": "postgresql://u:newtoken@h:5432/db"} - } + assert captured_kwargs == {"datasource": {"url": "postgresql://u:newtoken@h:5432/db"}} @pytest.mark.asyncio @@ -820,16 +796,12 @@ async def test_writer_recreate_does_not_use_datasource(monkeypatch): iam_token_db_auth=True, ) - await writer.recreate_prisma_client( - "postgresql://u:newtoken@h:5432/db", http_client=None - ) + await writer.recreate_prisma_client("postgresql://u:newtoken@h:5432/db", http_client=None) assert "datasource" not in captured_kwargs -def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( - monkeypatch, caplog -): +def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails(monkeypatch, caplog): """A transient AWS STS error (or any other failure) during the reader IAM token mint must NOT abort proxy startup. The reader is opt-in, so `PrismaClient.__init__` should log a warning and fall back to the @@ -862,9 +834,7 @@ def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( raise RuntimeError("simulated AWS STS hiccup") fake_iam_module.generate_iam_auth_token = boom - monkeypatch.setitem( - sys.modules, "litellm.proxy.auth.rds_iam_token", fake_iam_module - ) + monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_iam_module) from litellm.proxy.utils import PrismaClient @@ -879,9 +849,45 @@ def test_prisma_client_init_falls_back_to_writer_when_reader_iam_token_fails( assert isinstance(client.db, PrismaWrapper) assert not isinstance(client.db, RoutingPrismaWrapper) # And the operator gets a clear warning. - assert any( - "Failed to initialize read replica Prisma client" in r.getMessage() - for r in caplog.records + assert any("Failed to initialize read replica Prisma client" in r.getMessage() for r in caplog.records) + + +def test_prisma_client_configures_azure_token_refresh_for_writer_and_reader(monkeypatch): + from litellm.proxy.db.prisma_client import DatabaseTokenAuth + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False) + monkeypatch.setenv("AZURE_POSTGRESQL_AUTH", "true") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", + "postgresql://reader%40example.com@reader.postgres.database.azure.com:5432/litellm", + ) + + class FakePrisma: + def __init__(self, **kwargs): + self.kwargs = kwargs + + fake_prisma_module = MagicMock() + fake_prisma_module.Prisma = FakePrisma + monkeypatch.setitem(sys.modules, "prisma", fake_prisma_module) + + with patch( + "litellm.proxy.auth.azure_postgres_token.generate_azure_postgres_auth_token", + return_value="AZURE_TOKEN", + ): + from litellm.proxy.utils import PrismaClient + + client = PrismaClient( + database_url="postgresql://writer@writer.postgres.database.azure.com:5432/litellm", + proxy_logging_obj=MagicMock(), + ) + + assert isinstance(client.db, RoutingPrismaWrapper) + assert client.db.writer.database_token_auth == DatabaseTokenAuth.AZURE_ENTRA + assert client.db.reader.database_token_auth == DatabaseTokenAuth.AZURE_ENTRA + assert ( + os.environ["DATABASE_URL_READ_REPLICA"] + == "postgresql://reader%40example.com:AZURE_TOKEN@reader.postgres.database.azure.com:5432/litellm" ) @@ -940,10 +946,7 @@ async def test_connect_logs_writer_degradation(caplog): with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await routing.connect() - assert any( - "Failed to connect to primary (writer) DB" in r.getMessage() - for r in caplog.records - ) + assert any("Failed to connect to primary (writer) DB" in r.getMessage() for r in caplog.records) @pytest.mark.asyncio @@ -974,9 +977,7 @@ async def test_recreate_keeps_writer_unavailable_when_writer_recreate_fails(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper writer = MagicMock() - writer.recreate_prisma_client = AsyncMock( - side_effect=RuntimeError("primary still down") - ) + writer.recreate_prisma_client = AsyncMock(side_effect=RuntimeError("primary still down")) reader = MagicMock() reader.iam_token_db_auth = False reader.recreate_prisma_client = AsyncMock() diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py index f02f59cccc0..cf1bcd44da6 100644 --- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py @@ -1,20 +1,40 @@ -import json import os import sys -from typing import Optional from unittest.mock import MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules sys.path.insert(0, os.path.abspath("../..")) -import pytest - from litellm.secret_managers.get_azure_ad_token_provider import ( + build_azure_identity_credential, get_azure_ad_token_provider, ) class TestGetAzureAdTokenProvider: + @patch.dict( + os.environ, + { + "AZURE_CLIENT_ID": "test-client-id", + "AZURE_TENANT_ID": "test-tenant-id", + "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token", + }, + clear=True, + ) + @patch("azure.identity.WorkloadIdentityCredential") + def test_build_azure_identity_credential_uses_workload_identity( + self, + mock_workload_identity_credential, + ): + credential = build_azure_identity_credential() + + assert credential is mock_workload_identity_credential.return_value + mock_workload_identity_credential.assert_called_once_with( + client_id="test-client-id", + tenant_id="test-tenant-id", + token_file_path="/var/run/secrets/azure/tokens/azure-identity-token", + ) + @patch.dict( os.environ, { @@ -84,9 +104,7 @@ class TestGetAzureAdTokenProvider: # Assertions assert callable(result) - mock_managed_identity_credential.assert_called_once_with( - client_id="test-client-id" - ) + mock_managed_identity_credential.assert_called_once_with(client_id="test-client-id") mock_get_bearer_token_provider.assert_called_once_with( mock_credential_instance, "https://cognitiveservices.azure.com/.default" ) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 4b9f13c340b..71a33d66f09 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -14,6 +14,7 @@ from litellm._redis import ( ) from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL from litellm._redis_credential_provider import ( + AzureADCredentialProvider, GCPIAMCredentialProvider, _token_cache, ) @@ -27,6 +28,57 @@ def clear_gcp_iam_token_cache(): _token_cache.clear() +def test_sync_redis_url_passes_azure_credential_provider(): + credential = MagicMock() + credential.get_token.return_value.token = "azure-token" + + with ( + patch("litellm._redis._build_azure_credential", return_value=credential), + patch("litellm._redis.redis.Redis.from_url") as from_url, + ): + get_redis_client( + url="rediss://cache.redis.cache.windows.net:6380", + azure_redis_ad_token="true", + username="entra-user", + ) + + provider = from_url.call_args.kwargs["credential_provider"] + assert isinstance(provider, AzureADCredentialProvider) + assert provider.get_credentials() == ("entra-user", "azure-token") + + +def test_async_redis_url_passes_azure_credential_provider(): + credential = MagicMock() + + with ( + patch("litellm._redis._build_azure_credential", return_value=credential), + patch("litellm._redis.async_redis.Redis.from_url") as from_url, + ): + get_redis_async_client( + url="rediss://cache.redis.cache.windows.net:6380", + azure_redis_ad_token="true", + username="entra-user", + ) + + assert isinstance(from_url.call_args.kwargs["credential_provider"], AzureADCredentialProvider) + + +def test_redis_url_connection_pool_passes_azure_credential_provider(): + credential = MagicMock() + + with ( + patch("litellm._redis._build_azure_credential", return_value=credential), + patch("litellm._redis.async_redis.BlockingConnectionPool.from_url") as from_url, + ): + get_redis_connection_pool( + url="rediss://cache.redis.cache.windows.net:6380", + azure_redis_ad_token="true", + username="entra-user", + ) + + assert isinstance(from_url.call_args.kwargs["credential_provider"], AzureADCredentialProvider) + + def test_get_redis_url_from_environment_single_url(monkeypatch): """Test when REDIS_URL is directly provided""" # Set the environment variable @@ -133,10 +185,7 @@ def test_get_redis_url_from_environment_missing_host_port(monkeypatch): get_redis_url_from_environment() # Check the error message - assert ( - "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" - in str(excinfo.value) - ) + assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) def test_get_redis_url_from_environment_missing_port(monkeypatch): @@ -151,18 +200,13 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): get_redis_url_from_environment() # Check the error message - assert ( - "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" - in str(excinfo.value) - ) + assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value) def test_max_connections_in_cluster_kwargs(): """Test that max_connections is included in Redis cluster kwargs""" kwargs = _get_redis_cluster_kwargs() - assert ( - "max_connections" in kwargs - ), "max_connections should be in available Redis cluster kwargs" + assert "max_connections" in kwargs, "max_connections should be in available Redis cluster kwargs" def test_socket_timeouts_in_cluster_kwargs(): @@ -222,7 +266,6 @@ def test_get_redis_async_client_with_connection_pool(): patch("litellm._redis.async_redis.Redis") as mock_redis, patch("litellm._redis._get_redis_client_logic") as mock_logic, ): - # Configure mock to return basic redis kwargs mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} @@ -231,12 +274,8 @@ def test_get_redis_async_client_with_connection_pool(): # Verify Redis was called with connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert ( - "connection_pool" in call_kwargs - ), "connection_pool should be passed to Redis client" - assert ( - call_kwargs["connection_pool"] == mock_pool - ), "connection_pool should match the provided pool" + assert "connection_pool" in call_kwargs, "connection_pool should be passed to Redis client" + assert call_kwargs["connection_pool"] == mock_pool, "connection_pool should match the provided pool" def test_get_redis_async_client_without_connection_pool(): @@ -245,7 +284,6 @@ def test_get_redis_async_client_without_connection_pool(): patch("litellm._redis.async_redis.Redis") as mock_redis, patch("litellm._redis._get_redis_client_logic") as mock_logic, ): - # Configure mock to return basic redis kwargs mock_logic.return_value = {"host": "localhost", "port": 6379, "db": 0} @@ -254,9 +292,7 @@ def test_get_redis_async_client_without_connection_pool(): # Verify Redis was called without connection_pool in kwargs call_kwargs = mock_redis.call_args[1] - assert ( - "connection_pool" not in call_kwargs - ), "connection_pool should not be in kwargs when not provided" + assert "connection_pool" not in call_kwargs, "connection_pool should not be in kwargs when not provided" def test_gcp_iam_credential_provider_get_credentials(): @@ -328,9 +364,7 @@ def test_gcp_iam_credential_provider_cache_shared_across_instances(): share one cached token so concurrent Redis connections don't each trigger a blocking IAM round-trip. """ - service_account = ( - "projects/-/serviceAccounts/shared@project.iam.gserviceaccount.com" - ) + service_account = "projects/-/serviceAccounts/shared@project.iam.gserviceaccount.com" with patch( "litellm._redis_credential_provider._generate_gcp_iam_access_token", @@ -354,35 +388,21 @@ def test_get_redis_async_client_gcp_cluster_uses_credential_provider(): """ startup_nodes = [{"host": "redis-node-1", "port": 6379}] - mock_connect_func = MagicMock() - mock_connect_func._gcp_service_account = ( - "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com" - ) - - redis_kwargs = { - "startup_nodes": startup_nodes, - "redis_connect_func": mock_connect_func, - } - - with ( - patch("litellm._redis.async_redis.RedisCluster") as mock_cluster, - patch("litellm._redis._get_redis_client_logic", return_value=redis_kwargs), - ): - get_redis_async_client() + with patch("litellm._redis.async_redis.RedisCluster") as mock_cluster: + get_redis_async_client( + startup_nodes=startup_nodes, + gcp_service_account="projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", + ) assert mock_cluster.called cluster_call_kwargs = mock_cluster.call_args[1] # Must use credential_provider, not a static password - assert ( - "credential_provider" in cluster_call_kwargs - ), "async GCP cluster must use credential_provider for per-connection token refresh" - assert isinstance( - cluster_call_kwargs["credential_provider"], GCPIAMCredentialProvider + assert "credential_provider" in cluster_call_kwargs, ( + "async GCP cluster must use credential_provider for per-connection token refresh" ) - assert ( - "password" not in cluster_call_kwargs - ), "async GCP cluster must not use a static password (expires after 1h)" + assert isinstance(cluster_call_kwargs["credential_provider"], GCPIAMCredentialProvider) + assert "password" not in cluster_call_kwargs, "async GCP cluster must not use a static password (expires after 1h)" @patch("litellm._redis.init_redis_cluster") @@ -399,9 +419,7 @@ def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch): mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert ( - "startup_nodes" in call_kwargs - ), "startup_nodes must be forwarded to init_redis_cluster" + assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to init_redis_cluster" @patch("litellm._redis.async_redis.RedisCluster") @@ -417,18 +435,12 @@ def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch): mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert ( - "startup_nodes" in call_kwargs - ), "startup_nodes must be forwarded to async RedisCluster" - assert ( - len(call_kwargs["startup_nodes"]) == 1 - ), "should forward exactly 1 cluster node" + assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" + assert len(call_kwargs["startup_nodes"]) == 1, "should forward exactly 1 cluster node" @patch("litellm._redis.async_redis.RedisCluster") -def test_async_client_prefers_cluster_over_url_via_env_var( - mock_cluster_cls, monkeypatch -): +def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, monkeypatch): """ Test get_redis_async_client returns async RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -443,15 +455,11 @@ def test_async_client_prefers_cluster_over_url_via_env_var( mock_cluster_cls.assert_called_once() call_kwargs = mock_cluster_cls.call_args[1] - assert ( - "startup_nodes" in call_kwargs - ), "startup_nodes must be forwarded to async RedisCluster" + assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster" @patch("litellm._redis.init_redis_cluster") -def test_sync_client_prefers_cluster_over_url_via_env_var( - mock_init_cluster, monkeypatch -): +def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, monkeypatch): """ Test get_redis_client returns RedisCluster when REDIS_CLUSTER_NODES is set even if REDIS_URL is also set. @@ -467,9 +475,7 @@ def test_sync_client_prefers_cluster_over_url_via_env_var( mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert ( - "startup_nodes" in call_kwargs - ), "startup_nodes must be forwarded to init_redis_cluster" + assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to init_redis_cluster" assert len(call_kwargs["startup_nodes"]) == 1 @@ -588,9 +594,7 @@ def test_async_sentinel_uses_sentinel_password_and_master_password( @patch("litellm._redis.init_redis_cluster") -def test_sync_client_preserves_password_for_cluster_when_url_also_set( - mock_init_cluster, monkeypatch -): +def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_cluster, monkeypatch): """ Test _get_redis_client_logic does not strip password from redis_kwargs when startup_nodes is present even if REDIS_URL is also set. @@ -604,9 +608,7 @@ def test_sync_client_preserves_password_for_cluster_when_url_also_set( mock_init_cluster.assert_called_once() call_kwargs = mock_init_cluster.call_args[0][0] - assert ( - "password" in call_kwargs - ), "password must not be stripped when routing to cluster" + assert "password" in call_kwargs, "password must not be stripped when routing to cluster" assert call_kwargs["password"] == "secret"