scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(proxy): Azure Entra ID DB auth (IAM_DB_AUTH_PROVIDER=azure)

Browse source
This commit is contained in:
mkceichelbeck 2026-06-05 12:05:37 +02:00
parent e9417603a3
commit 954412d5fc
11 changed files with 692 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/proxy/auth/azure_entra_db_token.py Normal file
View file

@ -0,0 +1,72 @@
"""Mint Microsoft Entra ID access tokens for Azure Database for PostgreSQL.
Azure counterpart to ``rds_iam_token.py``. Whereas AWS RDS mints a
host/port/user-scoped presigned IAM token, Azure issues a generic Microsoft
Entra ID access token for the Azure Database for PostgreSQL data-plane scope.
The access token is used as the Postgres *password*; the database username is
the Entra principal (managed identity / service principal) name.
Auth is resolved by ``azure-identity``'s ``DefaultAzureCredential``, so this
works keyless on AKS Microsoft Entra Workload Identity (the projected
federated token at ``AZURE_FEDERATED_TOKEN_FILE``), on Azure VMs / Container
Apps via managed identity, and locally via the Azure CLI — no static secret
required. The credential object is built once and reused so the Azure SDK's
internal token cache + silent refresh apply across calls.
"""
import threading
import urllib.parse
from typing import Any, Optional
# Microsoft Entra ID scope for the Azure Database for PostgreSQL / MySQL data
# plane. The issued access token is presented to Postgres as the password.
AZURE_DB_SCOPE = "https://ossrdbms-aad.database.windows.net/.default"
# Built once and reused: DefaultAzureCredential caches tokens internally and
# refreshes them silently, so a single instance avoids rebuilding the
# credential chain on every proactive token refresh. The lock guards the lazy
# init: the writer and reader refresh loops can mint concurrently, and the mint
# runs in a worker thread (via asyncio.to_thread), so two threads could
# otherwise build (and one silently discard) separate credentials.
_cached_credential: Optional[Any] = None
_credential_lock = threading.Lock()
def _get_default_credential() -> Any:
global _cached_credential
if _cached_credential is None:
with _credential_lock:
if _cached_credential is None:
try:
from azure.identity import DefaultAzureCredential
except ImportError as e:
raise ImportError(
"azure-identity is required for Azure Database for "
"PostgreSQL identity-token auth. Install it with: "
"pip install azure-identity"
) from e
_cached_credential = DefaultAzureCredential()
return _cached_credential
def generate_azure_entra_db_token(
db_host: Optional[str] = None,
db_port: Optional[str] = None,
db_user: Optional[str] = None,
credential: Optional[Any] = None,
) -> str:
"""Return a URL-encoded Entra ID access token for use as the DB password.
``db_host`` / ``db_port`` / ``db_user`` are accepted for signature parity
with ``rds_iam_token.generate_iam_auth_token`` but are not required: unlike
an AWS RDS presigned token, an Entra access token is scoped to the identity
and the OSS RDBMS audience, not to a specific host. ``credential`` may be
injected (primarily for testing); otherwise a cached
``DefaultAzureCredential`` is used.
"""
cred = credential if credential is not None else _get_default_credential()
access_token = cred.get_token(AZURE_DB_SCOPE).token
# Quote for parity/safety with the AWS path. Entra tokens are JWTs using the
# URL-safe base64 alphabet, so quoting is effectively a no-op here, but it
# guarantees the token can never corrupt the assembled postgresql:// URL.
return urllib.parse.quote(access_token, safe="")

100
litellm/proxy/db/db_iam_token.py Normal file
View file

@ -0,0 +1,100 @@
"""Provider dispatch for database identity-token authentication.
LiteLLM's proxy can authenticate to its Postgres database with a short-lived,
auto-refreshed identity token instead of a static password (gated on the
``IAM_TOKEN_DB_AUTH`` flag). Two token providers are supported:
* ``aws`` — AWS RDS / Aurora IAM auth
(:mod:`litellm.proxy.auth.rds_iam_token`)
* ``azure`` — Microsoft Entra ID for Azure Database for PostgreSQL
(:mod:`litellm.proxy.auth.azure_entra_db_token`)
The provider is auto-detected from the database host — Azure Database for
PostgreSQL endpoints always end in ``.postgres.database.azure.com``; anything
else (e.g. AWS RDS/Aurora) uses the default AWS RDS IAM provider. Every
DB-token mint site (CLI startup, the componentized ``DatabaseURLSettings``
entrypoints, and the runtime ``PrismaWrapper`` refresh loop) routes through
:func:`generate_db_iam_token`, so the same provider is used consistently on
first connect and on every refresh.
"""
import os
import urllib.parse
from typing import Optional
DB_IAM_PROVIDER_AWS = "aws"
DB_IAM_PROVIDER_AZURE = "azure"
# Azure Database for PostgreSQL Flexible Server endpoints always end in this
# suffix (private-endpoint FQDNs included), so the host unambiguously
# identifies the cloud — no extra env var needed to select the provider.
_AZURE_DB_HOST_SUFFIX = ".postgres.database.azure.com"
def get_db_iam_auth_provider(db_host: Optional[str] = None) -> str:
"""Detect the DB IAM token provider from the database host.
Returns ``azure`` for an Azure Database for PostgreSQL endpoint
(``*.postgres.database.azure.com``), otherwise ``aws`` (the default — AWS
RDS / Aurora IAM). ``db_host`` falls back to the ``DATABASE_HOST`` env var
when not provided.
"""
host = (db_host or os.getenv("DATABASE_HOST") or "").lower()
if _AZURE_DB_HOST_SUFFIX in host:
return DB_IAM_PROVIDER_AZURE
return DB_IAM_PROVIDER_AWS
def generate_db_iam_token(
db_host: Optional[str] = None,
db_port: Optional[str] = None,
db_user: Optional[str] = None,
) -> str:
"""Mint a DB auth token for the detected provider.
Returns a URL-quoted token suitable for use as the Postgres password in a
``postgresql://user:<token>@host:port/name`` URL.
"""
if get_db_iam_auth_provider(db_host) == DB_IAM_PROVIDER_AZURE:
from litellm.proxy.auth.azure_entra_db_token import (
generate_azure_entra_db_token,
)
return generate_azure_entra_db_token(
db_host=db_host, db_port=db_port, db_user=db_user
)
# Default: AWS RDS / Aurora IAM auth.
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
return generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user)
def build_postgres_url(
*,
user: Optional[str],
token: str,
host: Optional[str],
port: str,
name: Optional[str],
schema: Optional[str] = None,
) -> str:
"""Assemble a ``postgresql://`` URL for IAM / identity-token DB auth.
Single source of truth for all DB-IAM URL assembly (CLI startup, the
``DatabaseURLSettings`` writer/reader, and the ``PrismaWrapper`` refresh),
so the encoding can't drift between connect and refresh, or between the
writer and reader endpoints.
URL-encodes the principal (``user``), database (``name``), and ``schema``,
so Azure Entra user principals (UPNs containing ``@``) and other reserved
characters can't corrupt the URL. This is a no-op for conventional AWS RDS
IAM identifiers. ``token`` is assumed already URL-encoded by the provider
mint.
"""
user_q = urllib.parse.quote(str(user), safe="")
name_q = urllib.parse.quote(str(name), safe="")
url = f"postgresql://{user_q}:{token}@{host}:{port}/{name_q}"
if schema:
url += f"?schema={urllib.parse.quote(schema, safe='')}"
return url

36
litellm/proxy/db/db_url_settings.py
View file

@ -37,9 +37,9 @@ from typing import Optional, 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
# Imported as a module (not `from ... import <fn>`) so the provider-specific
# token mint stays patchable at its canonical location in tests.
from litellm.proxy.db import db_iam_token
_IAM_ENV_KEY = "IAM_TOKEN_DB_AUTH"
_DEFAULT_PG_PORT = "5432"
@ -132,15 +132,17 @@ class DatabaseURLSettings(BaseSettings):
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(
token = db_iam_token.generate_db_iam_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
return db_iam_token.build_postgres_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:
@ -194,13 +196,17 @@ class DatabaseURLSettings(BaseSettings):
)
user = cast(str, user)
name = cast(str, name)
token = rds_iam_token.generate_iam_auth_token(
token = db_iam_token.generate_db_iam_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
return db_iam_token.build_postgres_url(
user=user,
token=token,
host=host,
port=port,
name=name,
schema=schema,
)
if user and name:
return self._password_url(

167
litellm/proxy/db/prisma_client.py
View file

@ -3,6 +3,8 @@ This file contains the PrismaWrapper class, which is used to wrap the Prisma cli
"""
import asyncio
import base64
import json
import os
import random
import signal
@ -20,11 +22,12 @@ from litellm.secret_managers.main import str_to_bool
@dataclass(frozen=True)
class IAMEndpoint:
"""Static parts of an RDS IAM-authenticated Postgres connection.
"""Static parts of an identity-token-authenticated Postgres connection.
The IAM token rotates every ~15 minutes; everything else (host, port, user,
database name, schema) stays fixed. We capture the static fields once so
refresh just regenerates the token and reassembles the URL.
The auth token rotates (RDS IAM ~15 min, Azure Entra ~60-90 min); everything
else (host, port, user, database name, schema) stays fixed. We capture the
static fields once — DECODED — so refresh just regenerates the token and
reassembles the URL via the shared builder (which re-encodes user/name).
"""
host: str
@ -34,10 +37,16 @@ class IAMEndpoint:
schema: Optional[str] = None
def build_url(self, token: str) -> str:
url = f"postgresql://{self.user}:{token}@{self.host}:{self.port}/{self.name}"
if self.schema:
url += f"?schema={self.schema}"
return url
from litellm.proxy.db.db_iam_token import build_postgres_url
return build_postgres_url(
user=self.user,
token=token,
host=self.host,
port=self.port,
name=self.name,
schema=self.schema,
)
def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint:
@ -59,26 +68,33 @@ def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint:
schema_vals = qs.get("schema")
if schema_vals:
schema = schema_vals[0]
# Store user + db name DECODED (urlparse does not percent-decode), so a
# reader URL with a raw or %-encoded UPN yields the bare principal here;
# build_url re-encodes it, keeping reader and writer URLs identical.
return IAMEndpoint(
host=parsed.hostname,
port=port,
user=parsed.username,
name=name,
user=urllib.parse.unquote(parsed.username),
name=urllib.parse.unquote(name),
schema=schema,
)
class PrismaWrapper:
"""
Wrapper around Prisma client that handles RDS IAM token authentication.
Wrapper around Prisma client that handles short-lived DB auth tokens.
When iam_token_db_auth is enabled, this wrapper:
1. Proactively refreshes IAM tokens before they expire (background task)
1. Proactively refreshes the token before it expires (background task)
2. Falls back to synchronous refresh if a token is found expired
3. Uses proper locking to prevent race conditions during reconnection
RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them
3 minutes before expiration to ensure uninterrupted database connectivity.
The token provider is detected from the database host: AWS RDS/Aurora IAM
tokens (valid ~15 min) or Microsoft Entra ID tokens for Azure Database for
PostgreSQL (valid ~60-90 min). The refresh is scheduled
from each token's parsed expiry (RDS presigned params or JWT ``exp``),
falling back to a fixed interval when expiry cannot be parsed. The buffer
below refreshes ahead of expiration to keep connectivity uninterrupted.
"""
# Buffer time in seconds before token expiration to trigger refresh
@ -171,7 +187,8 @@ class PrismaWrapper:
"""
Extract the token (password) from the DATABASE_URL.
The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters.
The token is either an AWS RDS presigned token (with X-Amz-Date /
X-Amz-Expires query params) or an Azure Entra ID JWT.
Important: We must parse the URL while it's still encoded to preserve structure,
then decode the password portion. Otherwise the '?' in the token breaks URL parsing.
@ -192,11 +209,24 @@ class PrismaWrapper:
"""
Parse the token to extract its expiration time.
Returns the datetime when the token expires, or None if parsing fails.
Handles both auth-token shapes the proxy can mint:
* AWS RDS IAM presigned tokens (``...?X-Amz-Date=...&X-Amz-Expires=...``)
* Azure Entra ID access tokens (JWTs carrying an ``exp`` claim)
Returns the datetime when the token expires, or None if parsing fails
(callers then fall back to ``FALLBACK_REFRESH_INTERVAL_SECONDS``).
"""
if token is None:
return None
aws_expiration = self._parse_aws_presigned_token_expiration(token)
if aws_expiration is not None:
return aws_expiration
return self._parse_jwt_token_expiration(token)
def _parse_aws_presigned_token_expiration(self, token: str) -> Optional[datetime]:
"""Parse expiry from an AWS RDS IAM presigned token's query params."""
try:
# Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&...
if "?" not in token:
@ -216,7 +246,35 @@ class PrismaWrapper:
return token_created + timedelta(seconds=expires_in)
except Exception as e:
verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}")
verbose_proxy_logger.debug(
f"Failed to parse AWS presigned token expiration: {e}"
)
return None
def _parse_jwt_token_expiration(self, token: str) -> Optional[datetime]:
"""Parse expiry from a JWT's ``exp`` claim (e.g. Azure Entra ID tokens).
Decodes the payload without signature verification — we only read the
public ``exp`` timestamp to schedule a proactive refresh, never to make
a trust decision.
"""
try:
parts = token.split(".")
if len(parts) != 3:
return None
payload_segment = parts[1]
# Restore base64url padding before decoding.
padding = "=" * (-len(payload_segment) % 4)
decoded = base64.urlsafe_b64decode(payload_segment + padding)
claims = json.loads(decoded)
exp = claims.get("exp")
if exp is None:
return None
# `exp` is a Unix timestamp (UTC); keep naive-UTC to match the rest
# of this class (which compares against datetime.utcnow()).
return datetime.utcfromtimestamp(int(exp))
except Exception as e:
verbose_proxy_logger.debug(f"Failed to parse JWT token expiration: {e}")
return None
def _calculate_seconds_until_refresh(self) -> float:
@ -273,41 +331,51 @@ class PrismaWrapper:
return datetime.utcnow() > expiration_time
def get_rds_iam_token(self) -> Optional[str]:
"""Generate a new RDS IAM token and update the configured DB URL env var.
"""Generate a new DB auth token and update the configured DB URL env var.
When the wrapper was constructed with an explicit `iam_endpoint`
(typical for a reader wrapper whose host/port/user came from a parsed
URL), use that. Otherwise fall back to the legacy DATABASE_HOST/PORT/
USER/NAME/SCHEMA env vars (writer behavior).
Mints an AWS RDS IAM token or an Azure Entra ID token depending on the
database host. When the wrapper was constructed with an explicit
`iam_endpoint` (typical for a reader wrapper whose host/port/user came
from a parsed URL), use that. Otherwise fall back to the legacy
DATABASE_HOST/PORT/USER/NAME/SCHEMA env vars (writer behavior).
"""
if not self.iam_token_db_auth:
return None
from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token
from litellm.proxy.db.db_iam_token import (
build_postgres_url,
generate_db_iam_token,
)
if self._iam_endpoint is not None:
endpoint = self._iam_endpoint
token = generate_iam_auth_token(
token = generate_db_iam_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.
# Default to the Postgres standard port; the AWS RDS token mint
# (botocore) embeds the literal string "None" in the presigned URL
# when the port is None, which then fails to parse. (The Azure
# provider ignores the port.)
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(
token = generate_db_iam_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}"
_db_url = build_postgres_url(
user=db_user,
token=token,
host=db_host,
port=db_port,
name=db_name,
schema=db_schema,
)
os.environ[self._db_url_env_var] = _db_url
return _db_url
@ -347,7 +415,7 @@ class PrismaWrapper:
"""
Start the background token refresh task.
This task proactively refreshes RDS IAM tokens before they expire,
This task proactively refreshes the DB auth token before it expires,
preventing connection failures. Should be called after the initial
Prisma client connection is established.
"""
@ -363,7 +431,7 @@ 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 DB auth token proactive refresh background task",
self._log_prefix,
)
@ -383,19 +451,19 @@ class PrismaWrapper:
pass
self._token_refresh_task = None
verbose_proxy_logger.info(
"%sStopped RDS IAM token refresh background task", self._log_prefix
"%sStopped DB auth token refresh background task", self._log_prefix
)
async def _token_refresh_loop(self) -> None:
"""
Background loop that proactively refreshes RDS IAM tokens before expiration.
Background loop that proactively refreshes the DB auth token before expiration.
Uses precise timing: calculates the exact sleep duration until the token
needs to be refreshed (expiration - 3 minute buffer), then refreshes.
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}DB auth token refresh loop started. "
f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration."
)
@ -406,25 +474,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}DB auth token 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
"%sProactively refreshing DB auth token...", self._log_prefix
)
await self._safe_refresh_token()
except asyncio.CancelledError:
verbose_proxy_logger.info(
"%sRDS IAM token refresh loop cancelled", self._log_prefix
"%sDB auth token refresh loop cancelled", self._log_prefix
)
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 DB auth token refresh loop: {e}. "
f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..."
)
# On error, wait before retrying to avoid tight error loops
@ -435,23 +503,26 @@ class PrismaWrapper:
async def _safe_refresh_token(self) -> None:
"""
Refresh the RDS IAM token with proper locking to prevent race conditions.
Refresh the DB auth token with proper locking to prevent race conditions.
Uses an asyncio lock to ensure only one refresh operation happens at a time,
preventing multiple concurrent reconnection attempts.
"""
async with self._reconnection_lock:
new_db_url = self.get_rds_iam_token()
# Mint + reassemble off the event loop: the provider token mint
# (boto3 sign / azure-identity get_token) is a blocking network call
# that would otherwise stall the loop (and /health/liveliness)
# during the periodic refresh.
new_db_url = await asyncio.to_thread(self.get_rds_iam_token)
if new_db_url:
await self.recreate_prisma_client(new_db_url)
self._last_refresh_time = datetime.utcnow()
verbose_proxy_logger.info(
"%sRDS IAM token refreshed successfully. New token valid for ~15 minutes.",
self._log_prefix,
"%sDB auth token refreshed successfully.", self._log_prefix
)
else:
verbose_proxy_logger.error(
"%sFailed to generate new RDS IAM token during proactive refresh",
"%sFailed to generate new DB auth token during proactive refresh",
self._log_prefix,
)
@ -489,7 +560,7 @@ class PrismaWrapper:
if running_loop is not None:
verbose_proxy_logger.warning(
"%sRDS IAM token expired in __getattr__ — proactive refresh "
"%sDB auth token 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,
@ -500,7 +571,7 @@ class PrismaWrapper:
running_loop.create_task(self._safe_refresh_token())
else:
verbose_proxy_logger.warning(
"%sRDS IAM token expired in __getattr__ — proactive refresh "
"%sDB auth token expired in __getattr__ — proactive refresh "
"may have failed. Triggering synchronous fallback refresh...",
self._log_prefix,
)
@ -514,7 +585,7 @@ class PrismaWrapper:
self._log_prefix,
)
else:
raise ValueError("Failed to get RDS IAM token")
raise ValueError("Failed to get DB auth token")
return original_attr

19
litellm/proxy/proxy_cli.py
View file

@ -930,7 +930,10 @@ def run_server( # noqa: PLR0915
### 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
from litellm.proxy.db.db_iam_token import (
build_postgres_url,
generate_db_iam_token,
)
db_host = os.getenv("DATABASE_HOST")
# Default to the Postgres standard port. Without a default,
@ -943,14 +946,18 @@ def run_server( # noqa: PLR0915
db_name = os.getenv("DATABASE_NAME")
db_schema = os.getenv("DATABASE_SCHEMA")
token = generate_iam_auth_token(
token = generate_db_iam_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}"
_db_url = build_postgres_url(
user=db_user,
token=token,
host=db_host,
port=db_port,
name=db_name,
schema=db_schema,
)
os.environ["DATABASE_URL"] = _db_url
os.environ["IAM_TOKEN_DB_AUTH"] = "True"

6
litellm/proxy/utils.py
View file

@ -2846,11 +2846,11 @@ class PrismaClient:
# `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,
from litellm.proxy.db.db_iam_token import (
generate_db_iam_token,
)
reader_token = generate_iam_auth_token(
reader_token = generate_db_iam_token(
db_host=reader_iam_endpoint.host,
db_port=reader_iam_endpoint.port,
db_user=reader_iam_endpoint.user,

119
tests/test_litellm/proxy/db/test_azure_db_url_encoding.py Normal file
View file

@ -0,0 +1,119 @@
"""DB-IAM URL assembly: principal/db-name/schema encoding and writer/reader parity.
Azure Entra **user** principals are UPNs containing '@', and managed-identity /
service-principal names can contain other reserved characters; embedding them
raw would corrupt the postgresql:// URL. Encoding is a no-op for simple AWS RDS
IAM usernames, so the AWS path is unaffected. The reader (parsed-URL) path must
produce the same byte-identical URL as the writer (env) path for the same
principal — the bug a single shared builder is meant to prevent.
Run:
uv run pytest tests/test_litellm/proxy/db/test_azure_db_url_encoding.py -v
"""
import os
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy.db import db_iam_token
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.prisma_client import PrismaWrapper, parse_iam_endpoint_from_url
_UPN = "svc@athenir.com"
_ENC = "svc%40athenir.com"
_HOST = "myserver.postgres.database.azure.com"
@pytest.fixture(autouse=True)
def _clean_env(monkeypatch):
monkeypatch.delenv("DATABASE_SCHEMA", raising=False)
yield
def _patch_token(value="TOK"):
return patch(
"litellm.proxy.db.db_iam_token.generate_db_iam_token", return_value=value
)
# --- principal encoding (writer / env paths) --------------------------------
def test_get_rds_iam_token_url_encodes_principal(monkeypatch):
monkeypatch.setenv("DATABASE_HOST", _HOST)
monkeypatch.setenv("DATABASE_PORT", "5432")
monkeypatch.setenv("DATABASE_USER", _UPN)
monkeypatch.setenv("DATABASE_NAME", "litellm")
wrapper = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=True)
with _patch_token("TOKEN123"):
url = wrapper.get_rds_iam_token()
assert url == f"postgresql://{_ENC}:TOKEN123@{_HOST}:5432/litellm"
assert f"{_UPN}:" not in url # raw '@' must not appear in the userinfo
assert os.environ["DATABASE_URL"] == url
def test_build_writer_url_encodes_principal(monkeypatch):
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv("DATABASE_HOST", _HOST)
monkeypatch.setenv("DATABASE_PORT", "5432")
monkeypatch.setenv("DATABASE_USER", _UPN)
monkeypatch.setenv("DATABASE_NAME", "litellm")
with _patch_token():
url = DatabaseURLSettings.from_env().build_writer_url()
assert url == f"postgresql://{_ENC}:TOK@{_HOST}:5432/litellm"
def test_simple_username_is_unchanged(monkeypatch):
"""Encoding is a no-op for simple AWS RDS IAM usernames (no regression)."""
monkeypatch.setenv("DATABASE_HOST", "db.rds.amazonaws.com")
monkeypatch.setenv("DATABASE_PORT", "5432")
monkeypatch.setenv("DATABASE_USER", "rds_iam_user")
monkeypatch.setenv("DATABASE_NAME", "litellm")
wrapper = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=True)
with _patch_token("T"):
url = wrapper.get_rds_iam_token()
assert url == "postgresql://rds_iam_user:T@db.rds.amazonaws.com:5432/litellm"
# --- build_postgres_url encoding --------------------------------------------
def test_build_postgres_url_basic():
url = db_iam_token.build_postgres_url(
user="u", token="t", host="h", port="5432", name="litellm"
)
assert url == "postgresql://u:t@h:5432/litellm"
def test_build_postgres_url_encodes_schema():
# A schema value containing a reserved char must be encoded, not injected
# raw as additional query parameters.
url = db_iam_token.build_postgres_url(
user="u", token="t", host="h", port="5432", name="db", schema="a&b"
)
assert url == "postgresql://u:t@h:5432/db?schema=a%26b"
# --- reader (parsed-URL) path parity ----------------------------------------
@pytest.mark.parametrize("reader_principal", [_UPN, _ENC])
def test_reader_endpoint_round_trip_matches_writer(reader_principal):
"""A reader URL written with either a raw OR a percent-encoded UPN parses to
the same decoded principal and rebuilds to the same encoded URL the writer
would emit (regression guard for the writer/reader divergence)."""
reader_url = (
f"postgresql://{reader_principal}:placeholder@reader.host:5432/litellm"
)
endpoint = parse_iam_endpoint_from_url(reader_url)
assert endpoint.user == _UPN # stored DECODED, regardless of input form
built = endpoint.build_url("FRESHTOKEN")
assert built == f"postgresql://{_ENC}:FRESHTOKEN@reader.host:5432/litellm"

93
tests/test_litellm/proxy/db/test_azure_entra_db_token.py Normal file
View file

@ -0,0 +1,93 @@
"""Unit tests for Azure Entra ID database token minting.
These tests inject a fake credential, so azure-identity does NOT need to be
installed to run them.
Run:
uv run pytest tests/test_litellm/proxy/db/test_azure_entra_db_token.py -v
"""
import urllib.parse
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from litellm.proxy.auth import azure_entra_db_token
from litellm.proxy.auth.azure_entra_db_token import (
AZURE_DB_SCOPE,
generate_azure_entra_db_token,
)
def _fake_credential(token_value: str) -> MagicMock:
cred = MagicMock()
cred.get_token.return_value = SimpleNamespace(token=token_value)
return cred
def test_uses_injected_credential_and_correct_scope():
cred = _fake_credential("header.payload.sig")
token = generate_azure_entra_db_token(credential=cred)
cred.get_token.assert_called_once_with(AZURE_DB_SCOPE)
# A JWT uses the url-safe base64 alphabet, so quoting is a no-op.
assert token == "header.payload.sig"
def test_url_quotes_reserved_characters():
# A token containing URL-reserved characters must be quoted so it cannot
# corrupt the assembled postgresql:// URL.
cred = _fake_credential("ab/cd=ef&gh")
token = generate_azure_entra_db_token(credential=cred)
assert token == urllib.parse.quote("ab/cd=ef&gh", safe="")
assert "/" not in token and "&" not in token and "=" not in token
def test_parity_args_do_not_affect_token():
cred = _fake_credential("tok")
token = generate_azure_entra_db_token(
db_host="h.postgres.database.azure.com",
db_port="5432",
db_user="mi-name",
credential=cred,
)
assert token == "tok"
def test_missing_azure_identity_raises_helpful_error(monkeypatch):
import builtins
real_import = builtins.__import__
def fake_import(name, *args, **kwargs):
if name == "azure.identity":
raise ImportError("simulated: azure-identity not installed")
return real_import(name, *args, **kwargs)
# Reset the module-level credential cache so the lazy import is attempted.
monkeypatch.setattr(azure_entra_db_token, "_cached_credential", None)
monkeypatch.setattr(builtins, "__import__", fake_import)
with pytest.raises(ImportError, match="azure-identity is required"):
generate_azure_entra_db_token()
def test_get_default_credential_builds_once_and_caches(monkeypatch):
"""The real (non-injected) path lazily builds DefaultAzureCredential exactly
once and reuses it (covers the credential cache + double-checked lock)."""
azure_identity = pytest.importorskip("azure.identity")
constructed = []
class _FakeCred:
def __init__(self):
constructed.append(1)
monkeypatch.setattr(azure_identity, "DefaultAzureCredential", _FakeCred)
monkeypatch.setattr(azure_entra_db_token, "_cached_credential", None)
first = azure_entra_db_token._get_default_credential()
second = azure_entra_db_token._get_default_credential()
assert first is second # cached
assert len(constructed) == 1 # built only once

76
tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py Normal file
View file

@ -0,0 +1,76 @@
"""Tests for Azure Entra ID (JWT) token-expiry parsing in PrismaWrapper.
The proactive refresh loop schedules the next refresh from the token's parsed
expiry. AWS RDS presigned tokens carry expiry in query params; Azure Entra ID
tokens are JWTs carrying an ``exp`` claim. Both must parse; anything else falls
back to the fixed refresh interval.
Run:
uv run pytest tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py -v
"""
import base64
import json
import time
from datetime import datetime, timedelta
from unittest.mock import MagicMock
from litellm.proxy.db.prisma_client import PrismaWrapper
def _b64url(data: dict) -> str:
raw = json.dumps(data).encode()
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
def _make_jwt(exp_epoch: int) -> str:
header = _b64url({"alg": "RS256", "typ": "JWT"})
payload = _b64url(
{"exp": exp_epoch, "aud": "https://ossrdbms-aad.database.windows.net"}
)
return f"{header}.{payload}.signature"
def _wrapper() -> PrismaWrapper:
return PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=True)
def test_parses_jwt_exp_claim():
exp = int(time.time()) + 3600
parsed = _wrapper()._parse_token_expiration(_make_jwt(exp))
assert parsed == datetime.utcfromtimestamp(exp)
def test_aws_presigned_token_still_parsed():
now = datetime.utcnow()
date_str = now.strftime("%Y%m%dT%H%M%SZ")
token = f"mock?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc"
parsed = _wrapper()._parse_token_expiration(token)
assert parsed is not None
assert abs((parsed - (now + timedelta(seconds=900))).total_seconds()) < 2
def test_unparseable_tokens_return_none():
w = _wrapper()
assert w._parse_token_expiration(None) is None
assert w._parse_token_expiration("not-a-token") is None
# A 3-segment string whose payload has no exp claim.
assert w._parse_token_expiration(_make_jwt_without_exp()) is None
def _make_jwt_without_exp() -> str:
header = _b64url({"alg": "RS256", "typ": "JWT"})
payload = _b64url({"aud": "https://ossrdbms-aad.database.windows.net"})
return f"{header}.{payload}.signature"
def test_extract_jwt_token_from_db_url_then_parse():
exp = int(time.time()) + 2700
jwt = _make_jwt(exp)
db_url = (
f"postgresql://mi-name:{jwt}@host.postgres.database.azure.com:5432/litellm"
)
w = _wrapper()
extracted = w._extract_token_from_db_url(db_url)
assert extracted == jwt
assert w._parse_token_expiration(extracted) == datetime.utcfromtimestamp(exp)

73
tests/test_litellm/proxy/db/test_db_iam_token.py Normal file
View file

@ -0,0 +1,73 @@
"""Unit tests for the DB IAM provider dispatcher (litellm.proxy.db.db_iam_token).
The provider is auto-detected from the database host: Azure Database for
PostgreSQL hosts (``*.postgres.database.azure.com``) use the Azure Entra path,
everything else uses the AWS RDS IAM default.
Run:
uv run pytest tests/test_litellm/proxy/db/test_db_iam_token.py -v
"""
from unittest.mock import patch
import pytest
from litellm.proxy.db import db_iam_token
@pytest.fixture(autouse=True)
def _clear_host_env(monkeypatch):
monkeypatch.delenv("DATABASE_HOST", raising=False)
yield
def test_default_provider_is_aws():
assert db_iam_token.get_db_iam_auth_provider() == "aws"
@pytest.mark.parametrize(
"host,expected",
[
("myserver.postgres.database.azure.com", "azure"),
("MYSERVER.POSTGRES.DATABASE.AZURE.COM", "azure"), # case-insensitive
("db.cluster-xyz.us-east-1.rds.amazonaws.com", "aws"),
("localhost", "aws"),
("", "aws"),
(None, "aws"),
],
)
def test_provider_detected_from_host(host, expected):
assert db_iam_token.get_db_iam_auth_provider(host) == expected
def test_provider_detected_from_database_host_env(monkeypatch):
monkeypatch.setenv("DATABASE_HOST", "x.postgres.database.azure.com")
assert db_iam_token.get_db_iam_auth_provider() == "azure"
def test_dispatch_aws_for_rds_host():
with patch(
"litellm.proxy.auth.rds_iam_token.generate_iam_auth_token",
return_value="aws-token",
) as aws_mock:
token = db_iam_token.generate_db_iam_token(
db_host="h.rds.amazonaws.com", db_port="5432", db_user="u"
)
assert token == "aws-token"
aws_mock.assert_called_once_with(
db_host="h.rds.amazonaws.com", db_port="5432", db_user="u"
)
def test_dispatch_azure_for_azure_host():
with patch(
"litellm.proxy.auth.azure_entra_db_token.generate_azure_entra_db_token",
return_value="azure-token",
) as az_mock:
token = db_iam_token.generate_db_iam_token(
db_host="h.postgres.database.azure.com", db_port="5432", db_user="mi-name"
)
assert token == "azure-token"
az_mock.assert_called_once_with(
db_host="h.postgres.database.azure.com", db_port="5432", db_user="mi-name"
)

6
tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py
View file

@ -553,18 +553,18 @@ async def test_iam_refresh_logs_carry_log_prefix(caplog):
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
await wrapper.start_token_refresh_task()
# Loop emits "RDS IAM token refresh loop started..." on first tick.
# Loop emits "DB auth token refresh loop started..." on first tick.
# Cancel immediately so the loop body runs once and we can assert.
await wrapper.stop_token_refresh_task()
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")
m.startswith("[reader] Started DB auth token proactive refresh")
for m in messages
)
assert any(
m.startswith("[reader] Stopped RDS IAM token refresh background task")
m.startswith("[reader] Stopped DB auth token refresh background task")
for m in messages
)