feat: add ability to auth to azure with token#21764
feat: add ability to auth to azure with token#21764Harshit28j wants to merge 1 commit intoBerriAI:mainfrom
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR adds Azure AD (Entra ID) passwordless authentication support for Azure Managed Redis caches, following the same pattern as the existing GCP IAM authentication. It introduces token generation via
Confidence Score: 3/5
|
| Filename | Overview |
|---|---|
| litellm/_redis.py | Adds Azure AD token generation, connect function, and integration into _get_redis_client_logic and get_redis_async_client for both standalone and cluster modes. Follows existing GCP IAM patterns well. Minor concern: no mutual exclusion between GCP IAM and Azure AD auth. |
| tests/test_litellm/test_utils.py | Adds Azure AD Redis tests. Mocking approach uses patch() on azure.identity module directly, which will fail in CI if azure-identity is not installed — unlike the GCP IAM tests that use patch.dict("sys.modules") for robustness. Also contains extensive formatting-only changes (90%+ of the diff). |
| docs/my-website/docs/caching/azure_redis_passwordless.md | New documentation page covering Azure Managed Redis passwordless authentication with Managed Identity, User-Assigned Identity, and Service Principal options. Clear and well-structured. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[_get_redis_client_logic] --> B{GCP IAM configured?}
B -->|Yes| C[create_gcp_iam_redis_connect_func]
B -->|No| D{Azure AD configured?}
C --> D
D -->|Yes| E[create_azure_ad_redis_connect_func]
D -->|No| F[Continue with standard auth]
E --> F
F --> G{Connection type?}
G -->|URL| H[Redis.from_url]
G -->|Cluster| I{Sync or Async?}
G -->|Sentinel| J[Redis Sentinel]
G -->|Standard| K[Redis client]
I -->|Sync| L[init_redis_cluster - passes redis_connect_func]
I -->|Async| M[get_redis_async_client]
M --> N{Has azure_redis_ad_token attr?}
N -->|Yes| O[_generate_azure_ad_redis_token]
O --> P[Set password + username on cluster]
N -->|No| Q{Has GCP service account?}
Q -->|Yes| R[_generate_gcp_iam_access_token]
R --> S[Set password on cluster]
Q -->|No| T[No IAM auth for cluster]
Last reviewed commit: 810ee68
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| mock_credential = Mock() | ||
| mock_credential.get_token.return_value = mock_token | ||
|
|
||
| with patch("azure.identity.DefaultAzureCredential", return_value=mock_credential): |
There was a problem hiding this comment.
Test will fail if azure-identity is not installed
The mock patch("azure.identity.DefaultAzureCredential", ...) requires the azure.identity module to be importable. If azure-identity is not installed in the CI environment, this patch will raise a ModuleNotFoundError before the test even runs.
Compare with the GCP IAM test (test_generate_gcp_iam_access_token) which uses patch.dict("sys.modules", {...}) to mock the entire module — that approach works regardless of whether the actual package is installed. The Azure tests should follow the same pattern for CI reliability.
For example:
mock_azure_identity = Mock()
mock_credential = Mock()
mock_credential.get_token.return_value = mock_token
mock_azure_identity.DefaultAzureCredential = Mock(return_value=mock_credential)
with patch.dict("sys.modules", {"azure.identity": mock_azure_identity}):
from litellm._redis import _generate_azure_ad_redis_token
result = _generate_azure_ad_redis_token()Context Used: Rule from dashboard - What: prevent any tests from being added here that make real network calls - only mock tests can be ... (source)
| mock_credential = Mock() | ||
| mock_credential.get_token.return_value = mock_token | ||
|
|
||
| with patch("azure.identity.ClientSecretCredential", return_value=mock_credential) as mock_cls: |
There was a problem hiding this comment.
Same mocking issue for ClientSecretCredential
Same as the DefaultAzureCredential test above — patch("azure.identity.ClientSecretCredential", ...) will fail if azure-identity is not installed. Use patch.dict("sys.modules", ...) instead, as was done in the GCP IAM tests.
Context Used: Rule from dashboard - What: prevent any tests from being added here that make real network calls - only mock tests can be ... (source)
| # 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": |
There was a problem hiding this comment.
Azure AD silently overwrites GCP IAM if both configured
If both gcp_service_account and azure_redis_ad_token: true are set, the Azure block (which runs second) will overwrite the redis_connect_func set by the GCP block without any warning. Consider adding mutual exclusion or at least a warning log to prevent silent misconfiguration.
1 participant
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
📖 Documentation
✅ Test
Changes
azure-identitylibrary to extract fresh OAuth access tokens mapped specifically forhttps://redis.azure.com/.defaultscopes.azure_redis_ad_token: true) mapped overazure_client_idspecifications.E2E Verification Details
azure-identitylibrary to return our local Redis password as a fake AD token.AUTH <token>command.x-litellm-cache-keyheader.