mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat: add ability to auth to azure with token
This commit is contained in:
Harshit Jain
Harshit28j
parent
4f4225fdbf
commit
3761598a02
3 changed files with 562 additions and 361 deletions
83
docs/my-website/docs/caching/azure_redis_passwordless.md
Normal file
83
docs/my-website/docs/caching/azure_redis_passwordless.md
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
---
|
||||
title: Azure Managed Redis Passwordless (IAM)
|
||||
---
|
||||
|
||||
# Azure Managed Redis Passwordless Authentication
|
||||
|
||||
LiteLLM supports [passwordless authentication to Azure Managed Redis](https://learn.microsoft.com/en-us/azure/azure-cache-for-redis/cache-azure-active-directory-for-authentication) using Azure Active Directory (Microsoft Entra ID). This allows you to securely connect to your Redis cache using Azure Managed Identities or Service Principals, avoiding the need to store static connection passwords.
|
||||
|
||||
## Prerequisites
|
||||
1. **Azure Cache for Redis** instance with **Azure AD Authentication enabled**.
|
||||
2. **`azure-identity` package** installed in your LiteLLM environment:
|
||||
```bash
|
||||
pip install azure-identity
|
||||
```
|
||||
3. Your Azure Identity must have a role assignment (e.g. `Data Owner` or `Data Contributor`) to the Redis cache.
|
||||
|
||||
## Configuration
|
||||
|
||||
Enable passwordless authentication in your LiteLLM configuration using `azure_redis_ad_token: true`.
|
||||
|
||||
### 1. Using System-Assigned Managed Identity
|
||||
|
||||
If running on an Azure service with a System-Assigned Managed Identity (e.g., Azure Container Apps, App Service, AKS), you don't need additional credentials. The `DefaultAzureCredential` will automatically discover the identity.
|
||||
|
||||
**Config (`config.yaml`)**:
|
||||
```yaml
|
||||
litellm_settings:
|
||||
cache: true
|
||||
cache_params:
|
||||
type: redis
|
||||
host: <my-redis-cache>.redis.cache.windows.net
|
||||
port: 6380
|
||||
ssl: true
|
||||
azure_redis_ad_token: true
|
||||
```
|
||||
|
||||
*Note: Azure Managed Redis mandates SSL, so `port: 6380` and `ssl: true` are required.*
|
||||
|
||||
### 2. Using User-Assigned Managed Identity
|
||||
|
||||
If using a User-Assigned Managed Identity, provide your `AZURE_CLIENT_ID` via environment variables.
|
||||
|
||||
**Environment variables**:
|
||||
```bash
|
||||
export AZURE_CLIENT_ID="<your-managed-identity-client-id>"
|
||||
# Optional. The Object ID of your identity. Defaults to empty string.
|
||||
export REDIS_USERNAME="<your-principal-object-id>"
|
||||
```
|
||||
|
||||
### 3. Using Service Principal
|
||||
|
||||
If authenticating via an Azure Service Principal, set the standard Azure identity environment variables:
|
||||
|
||||
**Environment variables**:
|
||||
```bash
|
||||
export AZURE_CLIENT_ID="<your-sp-client-id>"
|
||||
export AZURE_TENANT_ID="<your-sp-tenant-id>"
|
||||
export AZURE_CLIENT_SECRET="<your-sp-client-secret>"
|
||||
# Optional. The Object ID of your Service Principal. Defaults to empty string.
|
||||
export REDIS_USERNAME="<your-principal-object-id>"
|
||||
```
|
||||
|
||||
Alternatively, you can provide these directly in the `config.yaml`:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
cache: true
|
||||
cache_params:
|
||||
type: redis
|
||||
host: <my-redis-cache>.redis.cache.windows.net
|
||||
port: 6380
|
||||
ssl: true
|
||||
azure_redis_ad_token: true
|
||||
azure_client_id: os.environ/AZURE_CLIENT_ID
|
||||
azure_tenant_id: os.environ/AZURE_TENANT_ID
|
||||
azure_client_secret: os.environ/AZURE_CLIENT_SECRET
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
1. LiteLLM uses `azure-identity` to request short-lived access tokens explicitly scoped for Redis (`https://redis.azure.com/.default`).
|
||||
2. LiteLLM establishes a secure TLS connection with Redis and sends an `AUTH` command using the generated token.
|
||||
3. Every time the underlying `redis-py` connection disconnects or reconnects, LiteLLM intercepts the connection attempt to **generate a fresh token**, seamlessly handling token expiration.
|
||||
|
|
@ -34,7 +34,16 @@ def _get_redis_kwargs():
|
|||
"retry",
|
||||
}
|
||||
|
||||
include_args = ["url", "redis_connect_func", "gcp_service_account", "gcp_ssl_ca_certs"]
|
||||
include_args = [
|
||||
"url",
|
||||
"redis_connect_func",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_redis_ad_token",
|
||||
"azure_client_id",
|
||||
"azure_tenant_id",
|
||||
"azure_client_secret",
|
||||
]
|
||||
|
||||
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
|
||||
|
||||
|
|
@ -78,6 +87,10 @@ def _get_redis_cluster_kwargs(client=None):
|
|||
available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection
|
||||
available_args.append("gcp_service_account")
|
||||
available_args.append("gcp_ssl_ca_certs")
|
||||
available_args.append("azure_redis_ad_token")
|
||||
available_args.append("azure_client_id")
|
||||
available_args.append("azure_tenant_id")
|
||||
available_args.append("azure_client_secret")
|
||||
available_args.append("max_connections")
|
||||
|
||||
return available_args
|
||||
|
|
@ -103,10 +116,10 @@ def _redis_kwargs_from_environment():
|
|||
def _generate_gcp_iam_access_token(service_account: str) -> str:
|
||||
"""
|
||||
Generate GCP IAM access token for Redis authentication.
|
||||
|
||||
|
||||
Args:
|
||||
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
|
||||
|
||||
|
||||
Returns:
|
||||
Access token string for GCP IAM authentication
|
||||
"""
|
||||
|
|
@ -117,11 +130,11 @@ def _generate_gcp_iam_access_token(service_account: str) -> str:
|
|||
"google-cloud-iam is required for GCP IAM Redis authentication. "
|
||||
"Install it with: pip install google-cloud-iam"
|
||||
)
|
||||
|
||||
|
||||
client = iam_credentials_v1.IAMCredentialsClient()
|
||||
request = iam_credentials_v1.GenerateAccessTokenRequest(
|
||||
name=service_account,
|
||||
scope=['https://www.googleapis.com/auth/cloud-platform'],
|
||||
scope=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
response = client.generate_access_token(request=request)
|
||||
return str(response.access_token)
|
||||
|
|
@ -133,14 +146,15 @@ def create_gcp_iam_redis_connect_func(
|
|||
) -> Callable:
|
||||
"""
|
||||
Creates a custom Redis connection function for GCP IAM authentication.
|
||||
|
||||
|
||||
Args:
|
||||
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
|
||||
ssl_ca_certs: Path to SSL CA certificate file for secure connections
|
||||
|
||||
|
||||
Returns:
|
||||
A connection function that can be used with Redis clients
|
||||
"""
|
||||
|
||||
def iam_connect(self):
|
||||
"""Initialize the connection and authenticate using GCP IAM"""
|
||||
from redis.exceptions import (
|
||||
|
|
@ -148,55 +162,163 @@ def create_gcp_iam_redis_connect_func(
|
|||
AuthenticationWrongNumberOfArgsError,
|
||||
)
|
||||
from redis.utils import str_if_bytes
|
||||
|
||||
|
||||
self._parser.on_connect(self)
|
||||
|
||||
|
||||
auth_args = (_generate_gcp_iam_access_token(service_account),)
|
||||
self.send_command("AUTH", *auth_args, check_health=False)
|
||||
|
||||
|
||||
try:
|
||||
auth_response = self.read_response()
|
||||
except AuthenticationWrongNumberOfArgsError:
|
||||
# Fallback to password auth if IAM fails
|
||||
if hasattr(self, 'password') and self.password:
|
||||
if hasattr(self, "password") and self.password:
|
||||
self.send_command("AUTH", self.password, check_health=False)
|
||||
auth_response = self.read_response()
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
if str_if_bytes(auth_response) != "OK":
|
||||
raise AuthenticationError("GCP IAM authentication failed")
|
||||
|
||||
|
||||
return iam_connect
|
||||
|
||||
|
||||
def _generate_azure_ad_redis_token(
|
||||
azure_client_id: Optional[str] = None,
|
||||
azure_tenant_id: Optional[str] = None,
|
||||
azure_client_secret: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Generate Azure AD access token for Redis authentication.
|
||||
|
||||
Uses the azure-identity SDK to obtain a token scoped to Azure Cache for Redis.
|
||||
Supports:
|
||||
- Managed Identity (when running on Azure with AZURE_CLIENT_ID)
|
||||
- Service Principal (AZURE_CLIENT_ID + AZURE_TENANT_ID + AZURE_CLIENT_SECRET)
|
||||
- DefaultAzureCredential (automatic discovery)
|
||||
|
||||
Args:
|
||||
azure_client_id: Optional Azure client ID (overrides AZURE_CLIENT_ID env var)
|
||||
azure_tenant_id: Optional Azure tenant ID (overrides AZURE_TENANT_ID env var)
|
||||
azure_client_secret: Optional Azure client secret (overrides AZURE_CLIENT_SECRET env var)
|
||||
|
||||
Returns:
|
||||
Access token string for Azure Redis authentication
|
||||
"""
|
||||
try:
|
||||
from azure.identity import (
|
||||
ClientSecretCredential,
|
||||
DefaultAzureCredential,
|
||||
ManagedIdentityCredential,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"azure-identity is required for Azure AD Redis authentication. "
|
||||
"Install it with: pip install azure-identity"
|
||||
)
|
||||
|
||||
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
|
||||
|
||||
# Determine credential type
|
||||
_client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID")
|
||||
_tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID")
|
||||
_client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET")
|
||||
|
||||
if _client_id and _tenant_id and _client_secret:
|
||||
credential = ClientSecretCredential(
|
||||
client_id=_client_id,
|
||||
tenant_id=_tenant_id,
|
||||
client_secret=_client_secret,
|
||||
)
|
||||
elif _client_id:
|
||||
credential = ManagedIdentityCredential(client_id=_client_id)
|
||||
else:
|
||||
credential = DefaultAzureCredential()
|
||||
|
||||
token = credential.get_token(AZURE_REDIS_SCOPE)
|
||||
return token.token
|
||||
|
||||
|
||||
def create_azure_ad_redis_connect_func(
|
||||
azure_client_id: Optional[str] = None,
|
||||
azure_tenant_id: Optional[str] = None,
|
||||
azure_client_secret: Optional[str] = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
Creates a custom Redis connection function for Azure AD authentication.
|
||||
|
||||
Used for sync Redis clients. Generates a fresh Azure AD token on each
|
||||
connection/reconnection, ensuring token refresh is handled automatically.
|
||||
|
||||
Args:
|
||||
azure_client_id: Optional Azure client ID
|
||||
azure_tenant_id: Optional Azure tenant ID
|
||||
azure_client_secret: Optional Azure client secret
|
||||
|
||||
Returns:
|
||||
A connection function that can be used with Redis clients via `redis_connect_func`
|
||||
"""
|
||||
|
||||
def ad_connect(self):
|
||||
"""Initialize the connection and authenticate using Azure AD"""
|
||||
from redis.exceptions import (
|
||||
AuthenticationError,
|
||||
AuthenticationWrongNumberOfArgsError,
|
||||
)
|
||||
from redis.utils import str_if_bytes
|
||||
|
||||
self._parser.on_connect(self)
|
||||
|
||||
# Get username from REDIS_USERNAME env var or default principal ID
|
||||
username = os.environ.get("REDIS_USERNAME", "")
|
||||
|
||||
access_token = _generate_azure_ad_redis_token(
|
||||
azure_client_id=azure_client_id,
|
||||
azure_tenant_id=azure_tenant_id,
|
||||
azure_client_secret=azure_client_secret,
|
||||
)
|
||||
|
||||
# Azure Redis expects AUTH <username> <token> (Redis 6+ ACL style)
|
||||
auth_args = (username, access_token)
|
||||
self.send_command("AUTH", *auth_args, check_health=False)
|
||||
|
||||
try:
|
||||
auth_response = self.read_response()
|
||||
except AuthenticationWrongNumberOfArgsError:
|
||||
# Fallback: try with just the token
|
||||
self.send_command("AUTH", access_token, check_health=False)
|
||||
auth_response = self.read_response()
|
||||
|
||||
if str_if_bytes(auth_response) != "OK":
|
||||
raise AuthenticationError("Azure AD authentication failed for Redis")
|
||||
|
||||
return ad_connect
|
||||
|
||||
|
||||
def get_redis_url_from_environment():
|
||||
if "REDIS_URL" in os.environ:
|
||||
return os.environ["REDIS_URL"]
|
||||
|
||||
if "REDIS_HOST" not in os.environ or "REDIS_PORT" not in os.environ:
|
||||
raise ValueError(
|
||||
"Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis."
|
||||
)
|
||||
|
||||
raise ValueError("Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis.")
|
||||
|
||||
if "REDIS_SSL" in os.environ and os.environ["REDIS_SSL"].lower() == "true":
|
||||
redis_protocol = "rediss"
|
||||
else:
|
||||
redis_protocol = "redis"
|
||||
|
||||
|
||||
# Build authentication part of URL
|
||||
auth_part = ""
|
||||
if "REDIS_USERNAME" in os.environ and "REDIS_PASSWORD" in os.environ:
|
||||
auth_part = f"{os.environ['REDIS_USERNAME']}:{os.environ['REDIS_PASSWORD']}@"
|
||||
elif "REDIS_PASSWORD" in os.environ:
|
||||
auth_part = f"{os.environ['REDIS_PASSWORD']}@"
|
||||
|
||||
return (
|
||||
f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
|
||||
)
|
||||
|
||||
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
|
||||
|
||||
|
||||
def _get_redis_client_logic(**env_overrides):
|
||||
def _get_redis_client_logic(**env_overrides): # noqa: PLR0915
|
||||
"""
|
||||
Common functionality across sync + async redis client implementations
|
||||
"""
|
||||
|
|
@ -226,9 +348,9 @@ def _get_redis_client_logic(**env_overrides):
|
|||
if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str):
|
||||
redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes)
|
||||
|
||||
_sentinel_password: Optional[str] = redis_kwargs.get(
|
||||
"sentinel_password", None
|
||||
) or get_secret_str("REDIS_SENTINEL_PASSWORD")
|
||||
_sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str(
|
||||
"REDIS_SENTINEL_PASSWORD"
|
||||
)
|
||||
|
||||
if _sentinel_password is not None:
|
||||
redis_kwargs["sentinel_password"] = _sentinel_password
|
||||
|
|
@ -243,24 +365,49 @@ def _get_redis_client_logic(**env_overrides):
|
|||
# Handle GCP IAM authentication
|
||||
_gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
|
||||
_gcp_ssl_ca_certs = 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
|
||||
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
|
||||
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
# 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 = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
|
||||
|
||||
if _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true":
|
||||
_azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
|
||||
_azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
|
||||
_azure_client_secret = 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,
|
||||
)
|
||||
# Store Azure config on the function for async cluster access
|
||||
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
|
||||
redis_kwargs["redis_connect_func"]._azure_client_id = _azure_client_id
|
||||
redis_kwargs["redis_connect_func"]._azure_tenant_id = _azure_tenant_id
|
||||
redis_kwargs["redis_connect_func"]._azure_client_secret = _azure_client_secret
|
||||
|
||||
# Remove Azure-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("azure_redis_ad_token", None)
|
||||
redis_kwargs.pop("azure_client_id", None)
|
||||
redis_kwargs.pop("azure_tenant_id", None)
|
||||
redis_kwargs.pop("azure_client_secret", None)
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
redis_kwargs.pop("host", None)
|
||||
redis_kwargs.pop("port", None)
|
||||
|
|
@ -268,9 +415,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
redis_kwargs.pop("password", None)
|
||||
elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None:
|
||||
pass
|
||||
elif (
|
||||
"sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None
|
||||
):
|
||||
elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None:
|
||||
pass
|
||||
elif "host" not in redis_kwargs or redis_kwargs["host"] is None:
|
||||
raise ValueError("Either 'host' or 'url' must be specified for redis.")
|
||||
|
|
@ -313,9 +458,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
|
|||
service_name = redis_kwargs.get("service_name")
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError(
|
||||
"Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel."
|
||||
)
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
||||
verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.")
|
||||
|
||||
|
|
@ -337,9 +480,7 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
|||
service_name = redis_kwargs.get("service_name")
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError(
|
||||
"Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel."
|
||||
)
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
||||
verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.")
|
||||
|
||||
|
|
@ -376,8 +517,9 @@ def get_redis_client(**env_overrides):
|
|||
return redis.Redis(**redis_kwargs)
|
||||
|
||||
|
||||
def get_redis_async_client(
|
||||
connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **env_overrides,
|
||||
def get_redis_async_client( # noqa: PLR0915
|
||||
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
|
||||
**env_overrides,
|
||||
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
|
|
@ -390,9 +532,7 @@ def get_redis_async_client(
|
|||
url_kwargs[arg] = redis_kwargs[arg]
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(
|
||||
arg
|
||||
)
|
||||
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg)
|
||||
)
|
||||
return async_redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
|
|
@ -411,16 +551,20 @@ def get_redis_async_client(
|
|||
|
||||
# Get GCP service account - first try from redis_connect_func, then from environment
|
||||
gcp_service_account = None
|
||||
if redis_connect_func and hasattr(redis_connect_func, '_gcp_service_account'):
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
gcp_service_account = redis_connect_func._gcp_service_account
|
||||
else:
|
||||
gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
|
||||
|
||||
verbose_logger.debug(f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}"
|
||||
)
|
||||
|
||||
# If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password
|
||||
if redis_connect_func and gcp_service_account:
|
||||
verbose_logger.debug("DEBUG: Generating IAM token for service account (value not logged for security reasons)")
|
||||
verbose_logger.debug(
|
||||
"DEBUG: Generating IAM token for service account (value not logged for security reasons)"
|
||||
)
|
||||
try:
|
||||
# Generate IAM access token using the helper function
|
||||
access_token = _generate_gcp_iam_access_token(gcp_service_account)
|
||||
|
|
@ -429,21 +573,49 @@ def get_redis_async_client(
|
|||
except Exception as e:
|
||||
verbose_logger.error(f"Failed to generate GCP IAM access token: {e}")
|
||||
from redis.exceptions import AuthenticationError
|
||||
|
||||
raise AuthenticationError("Failed to generate GCP IAM access token")
|
||||
# Handle Azure AD authentication for async clusters
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_azure_redis_ad_token"):
|
||||
_az_client_id = getattr(redis_connect_func, "_azure_client_id", None)
|
||||
_az_tenant_id = getattr(redis_connect_func, "_azure_tenant_id", None)
|
||||
_az_client_secret = getattr(redis_connect_func, "_azure_client_secret", None)
|
||||
|
||||
verbose_logger.debug("Generating Azure AD token for async Redis cluster")
|
||||
try:
|
||||
access_token = _generate_azure_ad_redis_token(
|
||||
azure_client_id=_az_client_id,
|
||||
azure_tenant_id=_az_tenant_id,
|
||||
azure_client_secret=_az_client_secret,
|
||||
)
|
||||
cluster_kwargs["password"] = access_token
|
||||
# Set username if available
|
||||
_username = os.environ.get("REDIS_USERNAME", "")
|
||||
if _username:
|
||||
cluster_kwargs["username"] = _username
|
||||
verbose_logger.debug("Successfully generated Azure AD token for async Redis cluster")
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Failed to generate Azure AD access token: {e}")
|
||||
from redis.exceptions import AuthenticationError
|
||||
|
||||
raise AuthenticationError("Failed to generate Azure AD access token for Redis")
|
||||
else:
|
||||
verbose_logger.debug(f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
|
||||
|
||||
verbose_logger.debug(
|
||||
f"DEBUG: Not using GCP/Azure AD IAM auth - redis_connect_func={redis_connect_func is not None}"
|
||||
)
|
||||
|
||||
new_startup_nodes: List[ClusterNode] = []
|
||||
|
||||
for item in redis_kwargs["startup_nodes"]:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
cluster_kwargs.pop("startup_nodes", None)
|
||||
|
||||
|
||||
# Create async RedisCluster with IAM token as password if available
|
||||
cluster_client = async_redis.RedisCluster(
|
||||
startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore
|
||||
startup_nodes=new_startup_nodes,
|
||||
**cluster_kwargs, # type: ignore
|
||||
)
|
||||
|
||||
|
||||
return cluster_client
|
||||
|
||||
# Check for Redis Sentinel
|
||||
|
|
@ -463,7 +635,10 @@ def get_redis_connection_pool(**env_overrides):
|
|||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
pool_kwargs = {"timeout": REDIS_CONNECTION_POOL_TIMEOUT, "url": redis_kwargs["url"]}
|
||||
pool_kwargs = {
|
||||
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
|
||||
"url": redis_kwargs["url"],
|
||||
}
|
||||
if "max_connections" in redis_kwargs:
|
||||
try:
|
||||
pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])
|
||||
|
|
@ -479,9 +654,8 @@ def get_redis_connection_pool(**env_overrides):
|
|||
redis_kwargs.pop("ssl", None)
|
||||
redis_kwargs["connection_class"] = connection_class
|
||||
redis_kwargs.pop("startup_nodes", None)
|
||||
return async_redis.BlockingConnectionPool(
|
||||
timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs
|
||||
)
|
||||
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
|
||||
|
||||
|
||||
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
||||
"""Pretty print the Redis configuration using rich with sensitive data masking"""
|
||||
|
|
@ -492,6 +666,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
from rich.text import Text
|
||||
|
||||
if not verbose_logger.isEnabledFor(logging.DEBUG):
|
||||
return
|
||||
|
||||
|
|
@ -499,7 +674,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
|
||||
# Initialize the sensitive data masker
|
||||
masker = SensitiveDataMasker()
|
||||
|
||||
|
||||
# Mask sensitive data in redis_kwargs
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
|
||||
|
||||
|
|
@ -531,7 +706,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
value_str = str(value)
|
||||
else:
|
||||
value_str = str(value)
|
||||
|
||||
|
||||
config_table.add_row(key, value_str)
|
||||
|
||||
# Determine connection type
|
||||
|
|
@ -568,4 +743,3 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}")
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"Error pretty printing Redis configuration: {e}")
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue