fix(redis): type ElastiCache IAM configuration

Generated with AI

Co-Authored-By: Claude Code
This commit is contained in:
eugene-yao-zocdoc 2026-08-26 19:19:01 -04:00
parent 604b9edc53
commit f66891d80f
4 changed files with 65 additions and 51 deletions

View file

@ -279,25 +279,24 @@ def _is_true(value: object | None) -> bool:
return value is True or (isinstance(value, str) and value.lower() == "true")
def _build_elasticache_iam_provider(redis_kwargs: dict) -> ElastiCacheIAMCredentialProvider | None:
if not _is_true(redis_kwargs.get("aws_iam_auth")):
return None
required_settings: Final = {
"aws_iam_user_name": redis_kwargs.get("aws_iam_user_name"),
"aws_iam_cache_name": redis_kwargs.get("aws_iam_cache_name"),
"aws_iam_region": redis_kwargs.get("aws_iam_region")
or get_secret_str("AWS_REGION")
or get_secret_str("AWS_DEFAULT_REGION"),
}
missing_settings: Final = tuple(name for name, value in required_settings.items() if not value)
def _build_elasticache_iam_provider(
user_name: object | None,
cache_name: object | None,
region: object | None,
) -> ElastiCacheIAMCredentialProvider:
required_settings: Final = (
("aws_iam_user_name", user_name),
("aws_iam_cache_name", cache_name),
("aws_iam_region", region),
)
missing_settings: Final = tuple(name for name, value in required_settings if not value)
if missing_settings:
raise ValueError("AWS ElastiCache IAM Redis authentication requires: " + ", ".join(missing_settings))
return ElastiCacheIAMCredentialProvider(
user_name=str(required_settings["aws_iam_user_name"]),
cache_name=str(required_settings["aws_iam_cache_name"]),
region=str(required_settings["aws_iam_region"]),
user_name=str(user_name),
cache_name=str(cache_name),
region=str(region),
)
@ -605,10 +604,14 @@ def _get_redis_client_logic(**env_overrides):
"for Redis. Using Azure AD. Remove one to avoid misconfiguration."
)
elif _aws_iam_enabled:
aws_provider: Final = _build_elasticache_iam_provider(redis_kwargs)
if aws_provider is not None:
verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
redis_kwargs["credential_provider"] = aws_provider
verbose_logger.debug("Setting up AWS ElastiCache IAM authentication for Redis.")
redis_kwargs["credential_provider"] = _build_elasticache_iam_provider(
user_name=redis_kwargs.get("aws_iam_user_name"),
cache_name=redis_kwargs.get("aws_iam_cache_name"),
region=redis_kwargs.get("aws_iam_region")
or get_secret_str("AWS_REGION")
or get_secret_str("AWS_DEFAULT_REGION"),
)
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)

View file

@ -1,11 +1,19 @@
from __future__ import annotations
import asyncio
import threading
import time
from typing import Any, Final, Protocol
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Final, Protocol
from urllib.parse import quote
from redis.credentials import CredentialProvider
if TYPE_CHECKING:
from botocore.credentials import Credentials
else:
Credentials = Any # rebind-ok: runtime alias for the type-checking-only botocore import
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
@ -122,38 +130,24 @@ _ELASTICACHE_SERVICE_NAME: Final = "elasticache"
_ELASTICACHE_TOKEN_TTL_SECONDS: Final = 900
class _FrozenBotocoreCredentials(Protocol):
access_key: str
secret_key: str
token: str | None
class _BotocoreCredentials(Protocol):
def get_frozen_credentials(self) -> _FrozenBotocoreCredentials: ...
class _BotocoreCredentialsResolver(Protocol):
def __call__(self) -> _BotocoreCredentials | None: ...
class ElastiCacheIAMCredentialProvider(CredentialProvider):
def __init__(
self,
user_name: str,
cache_name: str,
region: str,
credentials_resolver: _BotocoreCredentialsResolver | None = None,
credentials_resolver: Callable[[], Credentials | None] | None = None,
token_lifetime_seconds: int = _ELASTICACHE_TOKEN_TTL_SECONDS,
) -> None:
self._user_name = user_name
self._cache_name = cache_name
self._region = region
self._credentials_resolver = credentials_resolver or self._resolve_credentials
self._credentials: _BotocoreCredentials | None = None
self._credentials: Credentials | None = None
self._token_lifetime_seconds = token_lifetime_seconds
@staticmethod
def _resolve_credentials() -> Any:
def _resolve_credentials() -> Credentials | None:
try:
import botocore.session
except ImportError as e:

View file

@ -4,10 +4,7 @@ from urllib.parse import parse_qs, urlsplit
import pytest
from litellm._redis_credential_provider import (
ElastiCacheIAMCredentialProvider,
_BotocoreCredentials,
)
from litellm._redis_credential_provider import ElastiCacheIAMCredentialProvider
class _FakeCredentials:
@ -20,16 +17,6 @@ class _FakeCredentials:
return self
class _FakeResolver:
def __init__(self, credentials: _BotocoreCredentials | None) -> None:
self.credentials = credentials
self.calls = 0
def __call__(self):
self.calls += 1
return self.credentials
class _RotatingFakeCredentials:
def __init__(self) -> None:
self.calls = 0
@ -46,6 +33,16 @@ class _RotatingFakeCredentials:
)
class _FakeResolver:
def __init__(self, credentials: _FakeCredentials | _RotatingFakeCredentials | None) -> None:
self.credentials = credentials
self.calls = 0
def __call__(self):
self.calls += 1
return self.credentials
def test_elasticache_provider_signs_expected_query():
resolver = _FakeResolver(_FakeCredentials("AKIA-SYNTHETIC"))
provider = ElastiCacheIAMCredentialProvider(

View file

@ -26357,6 +26357,26 @@ export interface components {
* independently of the response-cache backend in `litellm_settings.cache_params`.
*/
CoordinationRedisParams: {
/**
* Aws Iam Auth
* @description enable AWS ElastiCache IAM authentication
*/
aws_iam_auth?: boolean | string | null;
/**
* Aws Iam Cache Name
* @description AWS ElastiCache cache name
*/
aws_iam_cache_name?: string | null;
/**
* Aws Iam Region
* @description AWS region for ElastiCache IAM authentication
*/
aws_iam_region?: string | null;
/**
* Aws Iam User Name
* @description AWS ElastiCache IAM user name
*/
aws_iam_user_name?: string | null;
/**
* Host
* @description Redis hostname