feat(proxy): add Azure passwordless datastore auth

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

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 "" <token>
         # 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):

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
 

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="")

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

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:

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

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 ###
 

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

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

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)

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

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

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

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

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"