mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(redis): type ElastiCache IAM configuration
Generated with AI Co-Authored-By: Claude Code
This commit is contained in:
parent
604b9edc53
commit
f66891d80f
4 changed files with 65 additions and 51 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
20
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
20
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue