mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #38094 from eugene-yao-zocdoc/litellm_redis_credential_provider
fix(redis): support credential providers across clients
This commit is contained in:
commit
0ada822928
4 changed files with 540 additions and 67 deletions
|
|
@ -12,8 +12,9 @@ import json
|
|||
|
||||
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
|
|
@ -50,6 +51,7 @@ def _get_redis_kwargs():
|
|||
include_args: Final = {
|
||||
"url",
|
||||
"redis_connect_func",
|
||||
"credential_provider",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_redis_ad_token",
|
||||
|
|
@ -155,7 +157,8 @@ def _get_redis_cluster_kwargs(client=None):
|
|||
def _get_redis_env_kwarg_mapping():
|
||||
PREFIX: Final = "REDIS_"
|
||||
|
||||
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs()}
|
||||
exclude_from_environment: Final = frozenset({"credential_provider"})
|
||||
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment}
|
||||
|
||||
|
||||
def _redis_kwargs_from_environment():
|
||||
|
|
@ -353,6 +356,12 @@ def get_redis_url_from_environment():
|
|||
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
|
||||
|
||||
|
||||
def _url_without_userinfo(url: str) -> str:
|
||||
parts: Final = urlsplit(url)
|
||||
netloc: Final = parts.netloc.rsplit("@", 1)[-1]
|
||||
return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
|
||||
|
||||
|
||||
def _get_redis_client_logic(**env_overrides):
|
||||
"""
|
||||
Common functionality across sync + async redis client implementations
|
||||
|
|
@ -410,54 +419,58 @@ def _get_redis_client_logic(**env_overrides):
|
|||
if _service_name is not None:
|
||||
redis_kwargs["service_name"] = _service_name
|
||||
|
||||
# Handle GCP IAM authentication
|
||||
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
|
||||
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
|
||||
|
||||
if _gcp_service_account is not None:
|
||||
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
|
||||
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
|
||||
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
|
||||
if redis_kwargs.get("credential_provider") is None:
|
||||
# Handle GCP IAM authentication
|
||||
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str(
|
||||
"REDIS_GCP_SERVICE_ACCOUNT"
|
||||
)
|
||||
# Store GCP service account in redis_connect_func for async cluster access
|
||||
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
|
||||
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
|
||||
|
||||
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("gcp_service_account", None)
|
||||
redis_kwargs.pop("gcp_ssl_ca_certs", None)
|
||||
if _gcp_service_account is not None:
|
||||
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
|
||||
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
|
||||
|
||||
# Only enable SSL if explicitly requested AND SSL CA certs are provided
|
||||
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
|
||||
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
|
||||
# Only enable SSL if explicitly requested AND SSL CA certs are provided
|
||||
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
|
||||
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
|
||||
|
||||
# Handle Azure AD authentication (after GCP IAM block)
|
||||
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
|
||||
# Handle Azure AD authentication (after GCP IAM block)
|
||||
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
|
||||
|
||||
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
|
||||
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
|
||||
|
||||
if _azure_ad_enabled and _gcp_service_account is not None:
|
||||
verbose_logger.warning(
|
||||
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
|
||||
"Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
)
|
||||
if _azure_ad_enabled and _gcp_service_account is not None:
|
||||
verbose_logger.warning(
|
||||
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
|
||||
"Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
)
|
||||
|
||||
if _azure_ad_enabled and _gcp_service_account is None:
|
||||
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
|
||||
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
|
||||
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET")
|
||||
if _azure_ad_enabled and _gcp_service_account is None:
|
||||
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
|
||||
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
|
||||
_azure_client_secret: Final = 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_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
|
||||
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
|
||||
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
|
||||
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
|
||||
|
||||
redis_kwargs.pop("gcp_service_account", None)
|
||||
redis_kwargs.pop("gcp_ssl_ca_certs", None)
|
||||
|
||||
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("azure_redis_ad_token", None)
|
||||
|
|
@ -465,6 +478,13 @@ def _get_redis_client_logic(**env_overrides):
|
|||
redis_kwargs.pop("azure_tenant_id", None)
|
||||
redis_kwargs.pop("azure_client_secret", None)
|
||||
|
||||
if redis_kwargs.get("credential_provider") is not None:
|
||||
redis_kwargs.pop("redis_connect_func", None)
|
||||
redis_kwargs.pop("username", None)
|
||||
redis_kwargs.pop("password", None)
|
||||
if redis_kwargs.get("url") is not None:
|
||||
redis_kwargs["url"] = _url_without_userinfo(redis_kwargs["url"])
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
# Only strip host/port/db/password when not routing to a cluster.
|
||||
# When startup_nodes is also present the cluster path takes priority and
|
||||
|
|
@ -532,8 +552,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
|
|||
service_name: Final = redis_kwargs.get("service_name")
|
||||
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
|
||||
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
|
||||
sentinel_kwargs: Final = dict(connection_kwargs)
|
||||
sentinel_kwargs["password"] = sentinel_password
|
||||
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
|
@ -605,7 +624,12 @@ def _async_credential_provider(redis_connect_func: object | None) -> CredentialP
|
|||
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
|
||||
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
|
||||
which supersedes any static username or password redis-py would otherwise reject it with."""
|
||||
credential_provider: Final = _async_credential_provider(redis_kwargs.get("redis_connect_func"))
|
||||
explicit_provider: Final = redis_kwargs.get("credential_provider")
|
||||
credential_provider: Final = (
|
||||
explicit_provider
|
||||
if explicit_provider is not None
|
||||
else _async_credential_provider(redis_kwargs.get("redis_connect_func"))
|
||||
)
|
||||
if credential_provider is None:
|
||||
return redis_kwargs
|
||||
|
||||
|
|
@ -738,8 +762,20 @@ def get_redis_connection_pool(
|
|||
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
|
||||
|
||||
|
||||
def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {
|
||||
key: "<credential provider>"
|
||||
if key == "credential_provider" and value is not None
|
||||
else "<redis connect function>"
|
||||
if key == "redis_connect_func" and value is not None
|
||||
else value
|
||||
for key, value in redis_kwargs.items()
|
||||
}
|
||||
|
||||
|
||||
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
||||
"""Pretty print the Redis configuration using rich with sensitive data masking"""
|
||||
redis_kwargs_for_logging: Final = _redis_kwargs_for_logging(redis_kwargs)
|
||||
try:
|
||||
import logging
|
||||
|
||||
|
|
@ -757,7 +793,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
masker = SensitiveDataMasker()
|
||||
|
||||
# Mask sensitive data in redis_kwargs
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
|
||||
|
||||
# Create main panel title
|
||||
title: Final = Text("Redis Configuration", style="bold blue")
|
||||
|
|
@ -820,7 +856,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
except ImportError:
|
||||
# Fallback to simple logging if rich is not available
|
||||
masker = SensitiveDataMasker()
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
|
||||
verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.error("Error pretty printing Redis configuration: %s", e)
|
||||
|
|
|
|||
|
|
@ -175,6 +175,10 @@ _RedisCallResult = TypeVar("_RedisCallResult")
|
|||
_swallowed_redis_failures: Final[ContextVar[int]] = ContextVar("litellm_swallowed_redis_failures", default=0)
|
||||
|
||||
|
||||
def _opaque_kwarg_key(value: object) -> str:
|
||||
return f"{type(value).__name__}-{id(value)}"
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _redis_health_error_types() -> tuple[type, ...]:
|
||||
"""Exception types that mean the Redis backend itself is unhealthy.
|
||||
|
|
@ -399,10 +403,9 @@ class RedisCache(BaseCache):
|
|||
Generate a cache key for the async Redis client based on connection parameters.
|
||||
This ensures different Redis configurations use different cached clients.
|
||||
"""
|
||||
# Create a stable representation of redis_kwargs for hashing
|
||||
# Sort keys to ensure consistent hash regardless of parameter order
|
||||
sorted_kwargs: Final = sorted(self.redis_kwargs.items())
|
||||
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True)
|
||||
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True, default=_opaque_kwarg_key)
|
||||
kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16]
|
||||
return f"async-redis-client-{kwargs_hash}"
|
||||
|
||||
|
|
@ -1384,10 +1387,10 @@ class RedisCache(BaseCache):
|
|||
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
|
||||
"""
|
||||
try:
|
||||
import redis.asyncio as redis_async
|
||||
from .._redis import get_redis_async_client
|
||||
|
||||
# Create a fresh Redis client with current settings
|
||||
redis_client: Final = redis_async.Redis(**self.redis_kwargs)
|
||||
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
|
||||
|
||||
# Test the connection
|
||||
ping_result: Final = await redis_client.ping()
|
||||
|
|
|
|||
|
|
@ -64,22 +64,9 @@ class RedisClusterCache(RedisCache):
|
|||
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
|
||||
"""
|
||||
try:
|
||||
import redis.asyncio as redis_async
|
||||
from redis.cluster import ClusterNode
|
||||
from .._redis import get_redis_async_client
|
||||
|
||||
# Create ClusterNode objects from startup_nodes
|
||||
cluster_kwargs: Final = self.redis_kwargs.copy()
|
||||
startup_nodes: Final = cluster_kwargs.pop("startup_nodes", [])
|
||||
|
||||
new_startup_nodes: Final[list[ClusterNode]] = []
|
||||
for item in startup_nodes:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
|
||||
# Create a fresh Redis Cluster client with current settings
|
||||
redis_client: Final = redis_async.RedisCluster(
|
||||
startup_nodes=new_startup_nodes,
|
||||
**cluster_kwargs,
|
||||
)
|
||||
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
|
||||
|
||||
# Test the connection
|
||||
ping_result: Final = await redis_client.ping()
|
||||
|
|
|
|||
|
|
@ -1,13 +1,21 @@
|
|||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
from redis.credentials import CredentialProvider
|
||||
|
||||
import litellm
|
||||
from litellm._redis import (
|
||||
_async_auth_kwargs,
|
||||
_get_redis_client_logic,
|
||||
_get_redis_cluster_kwargs,
|
||||
_get_redis_env_kwarg_mapping,
|
||||
_get_redis_kwargs,
|
||||
_get_redis_url_kwargs,
|
||||
_pretty_print_redis_config,
|
||||
get_redis_async_client,
|
||||
get_redis_client,
|
||||
get_redis_connection_pool,
|
||||
|
|
@ -18,9 +26,69 @@ from litellm._redis_credential_provider import (
|
|||
GCPIAMCredentialProvider,
|
||||
_token_cache,
|
||||
)
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
|
||||
|
||||
|
||||
class _StubCredentialProvider(CredentialProvider):
|
||||
def __init__(self, token: str = "stub-token") -> None:
|
||||
self._token = token
|
||||
|
||||
def get_credentials(self):
|
||||
return (self._token,)
|
||||
|
||||
async def get_credentials_async(self):
|
||||
return (self._token,)
|
||||
|
||||
|
||||
class _HostileCredentialProvider(CredentialProvider):
|
||||
def __init__(self, secret: str) -> None:
|
||||
self._payload = secret
|
||||
|
||||
def get_credentials(self):
|
||||
return (self._payload,)
|
||||
|
||||
async def get_credentials_async(self):
|
||||
return (self._payload,)
|
||||
|
||||
def __repr__(self):
|
||||
raise AssertionError("provider repr must never be invoked")
|
||||
|
||||
def __str__(self):
|
||||
raise AssertionError("provider str must never be invoked")
|
||||
|
||||
def __reduce__(self):
|
||||
raise AssertionError("provider must never be serialized")
|
||||
|
||||
def __getstate__(self):
|
||||
raise AssertionError("provider state must never be inspected")
|
||||
|
||||
|
||||
def _gcp_marker_callback() -> MagicMock:
|
||||
callback = MagicMock()
|
||||
callback._gcp_service_account = "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"
|
||||
return callback
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_redis_environment(monkeypatch):
|
||||
for var in (
|
||||
"REDIS_URL",
|
||||
"REDIS_CLUSTER_NODES",
|
||||
"REDIS_SENTINEL_NODES",
|
||||
*_get_redis_env_kwarg_mapping(),
|
||||
):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clear_llm_client_cache():
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
yield
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_gcp_iam_token_cache():
|
||||
"""Reset the module-level GCP IAM token cache between tests."""
|
||||
|
|
@ -29,6 +97,364 @@ def clear_gcp_iam_token_cache():
|
|||
_token_cache.clear()
|
||||
|
||||
|
||||
def test_redis_allowlists_include_credential_provider():
|
||||
assert "credential_provider" in _get_redis_kwargs()
|
||||
assert "credential_provider" in _get_redis_url_kwargs()
|
||||
assert "credential_provider" in _get_redis_cluster_kwargs()
|
||||
|
||||
|
||||
def test_credential_provider_is_not_environment_derived():
|
||||
mapping = _get_redis_env_kwarg_mapping()
|
||||
assert "REDIS_CREDENTIAL_PROVIDER" not in mapping
|
||||
assert "credential_provider" not in mapping.values()
|
||||
|
||||
|
||||
def test_sync_direct_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_client(host="redis-host", port=6379, credential_provider=provider)
|
||||
|
||||
assert client.connection_pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_sync_direct_provider_supersedes_static_credentials(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_client(
|
||||
host="redis-host",
|
||||
port=6379,
|
||||
username="redis-user",
|
||||
password="redis-password",
|
||||
credential_provider=provider,
|
||||
)
|
||||
connection = client.connection_pool.make_connection()
|
||||
|
||||
assert connection.credential_provider is provider
|
||||
assert connection.username is None
|
||||
assert connection.password is None
|
||||
|
||||
|
||||
def test_sync_direct_provider_supersedes_environment_credentials(clean_redis_environment, monkeypatch):
|
||||
provider = _StubCredentialProvider()
|
||||
monkeypatch.setenv("REDIS_USERNAME", "redis-user")
|
||||
monkeypatch.setenv("REDIS_PASSWORD", "redis-password")
|
||||
|
||||
client = get_redis_client(host="redis-host", port=6379, credential_provider=provider)
|
||||
connection = client.connection_pool.make_connection()
|
||||
|
||||
assert connection.credential_provider is provider
|
||||
assert connection.username is None
|
||||
assert connection.password is None
|
||||
|
||||
|
||||
def test_sync_url_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_client(url="redis://redis-host:6379", credential_provider=provider)
|
||||
|
||||
assert client.connection_pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_async_direct_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_async_client(host="redis-host", port=6379, credential_provider=provider)
|
||||
|
||||
assert client.connection_pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_async_url_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_async_client(url="redis://redis-host:6379", credential_provider=provider)
|
||||
|
||||
assert client.connection_pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_sync_url_credentials_do_not_replace_explicit_provider(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_client(
|
||||
url="redis://url-user:url-pass@redis-host:6379",
|
||||
credential_provider=provider,
|
||||
)
|
||||
connection = client.connection_pool.make_connection()
|
||||
|
||||
assert connection.credential_provider is provider
|
||||
assert connection.username is None
|
||||
assert connection.password is None
|
||||
|
||||
|
||||
def test_async_url_credentials_do_not_replace_explicit_provider(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
client = get_redis_async_client(
|
||||
url="redis://url-user:url-pass@redis-host:6379",
|
||||
credential_provider=provider,
|
||||
)
|
||||
connection = client.connection_pool.make_connection()
|
||||
|
||||
assert connection.credential_provider is provider
|
||||
assert connection.username is None
|
||||
assert connection.password is None
|
||||
|
||||
|
||||
def test_async_host_port_pool_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
pool = get_redis_connection_pool(host="redis-host", port=6379, credential_provider=provider)
|
||||
|
||||
assert pool is not None
|
||||
assert pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_async_url_pool_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
pool = get_redis_connection_pool(url="redis://redis-host:6379", credential_provider=provider)
|
||||
|
||||
assert pool is not None
|
||||
assert pool.connection_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_async_url_pool_strips_userinfo_for_the_provider(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
pool = get_redis_connection_pool(url="rediss://url-user:url-pass@redis-host:6379/3", credential_provider=provider)
|
||||
|
||||
connection = pool.make_connection()
|
||||
assert connection.credential_provider is provider
|
||||
assert connection.username is None
|
||||
assert connection.password is None
|
||||
assert connection.db == 3
|
||||
|
||||
|
||||
def test_sync_cluster_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
startup_nodes = [{"host": "cluster-node", "port": 6379}]
|
||||
|
||||
with patch("redis.RedisCluster", autospec=True) as mock_cluster_cls:
|
||||
get_redis_client(startup_nodes=startup_nodes, credential_provider=provider, password="redis-secret")
|
||||
|
||||
cluster_kwargs = mock_cluster_cls.call_args.kwargs
|
||||
assert cluster_kwargs["credential_provider"] is provider
|
||||
assert "password" not in cluster_kwargs
|
||||
assert [(node.host, node.port) for node in cluster_kwargs["startup_nodes"]] == [("cluster-node", 6379)]
|
||||
|
||||
|
||||
def test_async_cluster_preserves_credential_provider_identity(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
startup_nodes = [{"host": "cluster-node", "port": 6379}]
|
||||
|
||||
client = get_redis_async_client(startup_nodes=startup_nodes, credential_provider=provider)
|
||||
|
||||
assert client.connection_kwargs["credential_provider"] is provider
|
||||
assert client.connection_kwargs["socket_keepalive"] is True
|
||||
assert client.connection_kwargs["health_check_interval"] == REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
|
||||
|
||||
|
||||
def test_explicit_provider_skips_automatic_auth_and_callback(clean_redis_environment, monkeypatch):
|
||||
provider = _StubCredentialProvider()
|
||||
monkeypatch.setenv("REDIS_GCP_SERVICE_ACCOUNT", "service-account@example.com")
|
||||
monkeypatch.setenv("REDIS_AZURE_AD_TOKEN", "true")
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: an auto-auth callback built here is popped again by the provider branch, so the builders are the only place the wasted work is visible
|
||||
"litellm._redis.create_gcp_iam_redis_connect_func"
|
||||
) as mock_gcp,
|
||||
patch( # test-quality-ok: same as above, and reaching this one also builds an Azure credential the caller never asked for
|
||||
"litellm._redis.create_azure_ad_redis_connect_func"
|
||||
) as mock_azure,
|
||||
):
|
||||
redis_kwargs = _get_redis_client_logic(
|
||||
host="redis-host",
|
||||
port=6379,
|
||||
credential_provider=provider,
|
||||
redis_connect_func=_gcp_marker_callback(),
|
||||
)
|
||||
|
||||
mock_gcp.assert_not_called()
|
||||
mock_azure.assert_not_called()
|
||||
assert redis_kwargs["credential_provider"] is provider
|
||||
assert "redis_connect_func" not in redis_kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
[
|
||||
{"gcp_ssl_ca_certs": "/tmp/ca.pem"},
|
||||
{"gcp_service_account": "sa@example.com", "gcp_ssl_ca_certs": "/tmp/ca.pem"},
|
||||
],
|
||||
ids=["certs-without-service-account", "both-alongside-a-provider"],
|
||||
)
|
||||
def test_gcp_kwargs_never_survive_client_logic(clean_redis_environment, overrides):
|
||||
redis_kwargs = _get_redis_client_logic(
|
||||
host="redis-host",
|
||||
port=6379,
|
||||
credential_provider=_StubCredentialProvider() if "gcp_service_account" in overrides else None,
|
||||
**overrides,
|
||||
)
|
||||
|
||||
assert "gcp_service_account" not in redis_kwargs
|
||||
assert "gcp_ssl_ca_certs" not in redis_kwargs
|
||||
|
||||
|
||||
def test_provider_keeps_the_rest_of_the_url_intact(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
redis_kwargs = _get_redis_client_logic(
|
||||
url="rediss://url-user:url-pass@redis-host:6379/3?protocol=3",
|
||||
credential_provider=provider,
|
||||
)
|
||||
|
||||
assert redis_kwargs["url"] == "rediss://redis-host:6379/3?protocol=3"
|
||||
|
||||
|
||||
def test_provider_free_url_is_left_untouched(clean_redis_environment):
|
||||
url = "redis://url-user:url-pass@redis-host:6379/3"
|
||||
|
||||
redis_kwargs = _get_redis_client_logic(url=url)
|
||||
|
||||
assert redis_kwargs["url"] == url
|
||||
|
||||
|
||||
def test_async_auth_kwargs_supersedes_credentials_an_explicit_provider_replaces():
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
auth_kwargs = _async_auth_kwargs(
|
||||
{
|
||||
"host": "redis-host",
|
||||
"port": 6379,
|
||||
"credential_provider": provider,
|
||||
"redis_connect_func": _gcp_marker_callback(),
|
||||
"username": "url-user",
|
||||
"password": "url-pass",
|
||||
}
|
||||
)
|
||||
|
||||
assert auth_kwargs["credential_provider"] is provider
|
||||
assert auth_kwargs["host"] == "redis-host"
|
||||
assert auth_kwargs["port"] == 6379
|
||||
assert "redis_connect_func" not in auth_kwargs
|
||||
assert "username" not in auth_kwargs
|
||||
assert "password" not in auth_kwargs
|
||||
|
||||
|
||||
def test_async_auth_kwargs_leaves_provider_free_kwargs_alone():
|
||||
redis_kwargs = {"host": "redis-host", "port": 6379, "username": "url-user", "password": "url-pass"}
|
||||
|
||||
assert _async_auth_kwargs(redis_kwargs) == redis_kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_cache_test_connection_uses_shared_factory(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
|
||||
with (
|
||||
patch("redis.Redis", autospec=True),
|
||||
patch("redis.asyncio.BlockingConnectionPool", autospec=True),
|
||||
patch("redis.asyncio.Redis", autospec=True) as mock_async_redis,
|
||||
):
|
||||
mock_async_redis.return_value.ping = AsyncMock(return_value=True)
|
||||
mock_async_redis.return_value.aclose = AsyncMock()
|
||||
cache = RedisCache(host="redis-host", port=6379, credential_provider=provider, password="redis-secret")
|
||||
result = await cache.test_connection()
|
||||
|
||||
client_kwargs = mock_async_redis.call_args.kwargs
|
||||
assert result["status"] == "success"
|
||||
assert client_kwargs["credential_provider"] is provider
|
||||
assert "password" not in client_kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_cluster_cache_test_connection_uses_shared_factory(clean_redis_environment):
|
||||
provider = _StubCredentialProvider()
|
||||
recorder = MagicMock()
|
||||
|
||||
class _StubAsyncCluster:
|
||||
def __init__(self, **kwargs):
|
||||
recorder(**kwargs)
|
||||
|
||||
async def ping(self):
|
||||
return True
|
||||
|
||||
async def aclose(self):
|
||||
return None
|
||||
|
||||
with (
|
||||
patch("redis.RedisCluster", autospec=True),
|
||||
patch("redis.asyncio.cluster.RedisCluster", _StubAsyncCluster),
|
||||
):
|
||||
cache = RedisClusterCache(startup_nodes=[{"host": "redis-host", "port": 6379}], credential_provider=provider)
|
||||
result = await cache.test_connection()
|
||||
|
||||
cluster_kwargs = recorder.call_args.kwargs
|
||||
assert result["status"] == "success"
|
||||
assert cluster_kwargs["credential_provider"] is provider
|
||||
|
||||
|
||||
def test_redis_cache_key_does_not_inspect_provider(clear_llm_client_cache):
|
||||
provider = _HostileCredentialProvider("synthetic-secret")
|
||||
second_provider = _StubCredentialProvider("another-token")
|
||||
|
||||
with (
|
||||
patch("redis.Redis", autospec=True),
|
||||
patch("redis.asyncio.BlockingConnectionPool", autospec=True),
|
||||
):
|
||||
cache = RedisCache(host="redis-host", port=6379, credential_provider=provider)
|
||||
second_cache = RedisCache(host="redis-host", port=6379, credential_provider=second_provider)
|
||||
|
||||
first_key = cache._get_async_client_cache_key()
|
||||
assert first_key == cache._get_async_client_cache_key()
|
||||
assert first_key != second_cache._get_async_client_cache_key()
|
||||
|
||||
|
||||
def test_pretty_print_never_expands_credential_provider(capsys):
|
||||
secret = "aaaa-UNIQUE-SENTINEL-bbbb"
|
||||
|
||||
with patch( # test-quality-ok: enable the debug-only printer without changing process-wide logger state
|
||||
"litellm._redis.verbose_logger.isEnabledFor", return_value=True
|
||||
):
|
||||
_pretty_print_redis_config(
|
||||
redis_kwargs={
|
||||
"host": "redis-host",
|
||||
"port": 6379,
|
||||
"credential_provider": _HostileCredentialProvider(secret),
|
||||
}
|
||||
)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert secret not in output
|
||||
assert "UNIQUE" not in output
|
||||
assert "_payload" not in output
|
||||
assert "credential_provider" in output
|
||||
|
||||
|
||||
def test_redis_cache_key_does_not_serialize_connect_func():
|
||||
def connect(connection):
|
||||
return None
|
||||
|
||||
cache = RedisCache.__new__(RedisCache)
|
||||
cache.redis_kwargs = {"host": "redis-host", "port": 6379, "redis_connect_func": connect}
|
||||
|
||||
first_key = cache._get_async_client_cache_key()
|
||||
assert first_key == cache._get_async_client_cache_key()
|
||||
|
||||
|
||||
def test_redis_cache_key_keys_opaque_kwargs_by_identity():
|
||||
|
||||
class _Opaque:
|
||||
pass
|
||||
|
||||
first = RedisCache.__new__(RedisCache)
|
||||
first.redis_kwargs = {"host": "redis-host", "retry": _Opaque()}
|
||||
second = RedisCache.__new__(RedisCache)
|
||||
second.redis_kwargs = {"host": "redis-host", "retry": _Opaque()}
|
||||
|
||||
assert first._get_async_client_cache_key() == first._get_async_client_cache_key()
|
||||
assert first._get_async_client_cache_key() != second._get_async_client_cache_key()
|
||||
|
||||
|
||||
def test_get_redis_url_from_environment_single_url(monkeypatch):
|
||||
"""Test when REDIS_URL is directly provided"""
|
||||
# Set the environment variable
|
||||
|
|
@ -500,6 +926,27 @@ def test_sync_sentinel_uses_sentinel_password_and_master_password(mock_sentinel_
|
|||
)
|
||||
|
||||
|
||||
@patch("redis.Sentinel")
|
||||
def test_sync_sentinel_keeps_provider_off_monitors_and_on_master(mock_sentinel_cls):
|
||||
provider = _StubCredentialProvider()
|
||||
mock_sentinel = MagicMock()
|
||||
mock_sentinel_cls.return_value = mock_sentinel
|
||||
|
||||
get_redis_client(
|
||||
sentinel_nodes=[("sentinel-1", 26379)],
|
||||
sentinel_password="sentinel-secret",
|
||||
service_name="mymaster",
|
||||
password="redis-secret",
|
||||
credential_provider=provider,
|
||||
)
|
||||
|
||||
sentinel_kwargs = mock_sentinel_cls.call_args.kwargs["sentinel_kwargs"]
|
||||
assert sentinel_kwargs["password"] == "sentinel-secret"
|
||||
assert "credential_provider" not in sentinel_kwargs
|
||||
assert mock_sentinel.master_for.call_args.kwargs["credential_provider"] is provider
|
||||
assert "password" not in mock_sentinel.master_for.call_args.kwargs
|
||||
|
||||
|
||||
@patch("litellm._redis.async_redis.Sentinel")
|
||||
def test_async_sentinel_uses_sentinel_password_and_master_password(
|
||||
mock_sentinel_cls,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue