From 954412d5fcedfaad3d1f75ff4ab46c39b507deb9 Mon Sep 17 00:00:00 2001 From: mkceichelbeck <141341133+mkceichelbeck@users.noreply.github.com> Date: Fri, 5 Jun 2026 12:05:37 +0200 Subject: [PATCH] feat(proxy): Azure Entra ID DB auth (IAM_DB_AUTH_PROVIDER=azure) --- litellm/proxy/auth/azure_entra_db_token.py | 72 ++++++++ litellm/proxy/db/db_iam_token.py | 100 +++++++++++ litellm/proxy/db/db_url_settings.py | 36 ++-- litellm/proxy/db/prisma_client.py | 167 +++++++++++++----- litellm/proxy/proxy_cli.py | 19 +- litellm/proxy/utils.py | 6 +- .../proxy/db/test_azure_db_url_encoding.py | 119 +++++++++++++ .../proxy/db/test_azure_entra_db_token.py | 93 ++++++++++ .../proxy/db/test_azure_entra_token_expiry.py | 76 ++++++++ .../proxy/db/test_db_iam_token.py | 73 ++++++++ .../proxy/db/test_routing_prisma_wrapper.py | 6 +- 11 files changed, 692 insertions(+), 75 deletions(-) create mode 100644 litellm/proxy/auth/azure_entra_db_token.py create mode 100644 litellm/proxy/db/db_iam_token.py create mode 100644 tests/test_litellm/proxy/db/test_azure_db_url_encoding.py create mode 100644 tests/test_litellm/proxy/db/test_azure_entra_db_token.py create mode 100644 tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py create mode 100644 tests/test_litellm/proxy/db/test_db_iam_token.py diff --git a/litellm/proxy/auth/azure_entra_db_token.py b/litellm/proxy/auth/azure_entra_db_token.py new file mode 100644 index 00000000000..d5f14ffe9a8 --- /dev/null +++ b/litellm/proxy/auth/azure_entra_db_token.py @@ -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="") diff --git a/litellm/proxy/db/db_iam_token.py b/litellm/proxy/db/db_iam_token.py new file mode 100644 index 00000000000..c1c04eb3179 --- /dev/null +++ b/litellm/proxy/db/db_iam_token.py @@ -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:@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 diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 58478db5e2e..dfc88485152 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -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 `) 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( diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index af5a58802bb..1e4e81e6de3 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -3,6 +3,8 @@ This file contains the PrismaWrapper class, which is used to wrap the Prisma cli """ import asyncio +import base64 +import 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 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c0246f234a8..406abac9738 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8bd50a50a38..fa8cf5364f4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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, diff --git a/tests/test_litellm/proxy/db/test_azure_db_url_encoding.py b/tests/test_litellm/proxy/db/test_azure_db_url_encoding.py new file mode 100644 index 00000000000..1e925bc3294 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_azure_db_url_encoding.py @@ -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" diff --git a/tests/test_litellm/proxy/db/test_azure_entra_db_token.py b/tests/test_litellm/proxy/db/test_azure_entra_db_token.py new file mode 100644 index 00000000000..7dfcb11f069 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_azure_entra_db_token.py @@ -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 diff --git a/tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py b/tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py new file mode 100644 index 00000000000..81207273ad4 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_azure_entra_token_expiry.py @@ -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) diff --git a/tests/test_litellm/proxy/db/test_db_iam_token.py b/tests/test_litellm/proxy/db/test_db_iam_token.py new file mode 100644 index 00000000000..9d3c6b65b4b --- /dev/null +++ b/tests/test_litellm/proxy/db/test_db_iam_token.py @@ -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" + ) diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index 8c3a2b9e2d7..d6231e31e9c 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -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 )