Skip to content

Commit 88f1047

Browse files
authored
Add native support for Azure DevOps OIDC authentication (#1027)
Added native support for Azure DevOps OIDC authentication to allow authentication from Azure DevOps pipeline's environment. ## Testing Tested on an Azure DevOps project. Used the SDK to authenticate with a Databrick's workspace using Azure DevOps OIDC 1) Created a demo python files that uses the SDK to create a Workspace Client: This will test if the OIDC authentication is working <img width="1905" height="937" alt="Screenshot 2025-09-18 at 15 42 18" src="https://github.com/user-attachments/assets/8fef9c12-a3a4-4421-b6e3-608e76072b21" /> 2) Created a pipeline to run the demo file. Set the necessary environment variables. (System.AccessToken is slightly different, it needs to be exported through pipeline syntax: https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken) <img width="1898" height="975" alt="Screenshot 2025-09-18 at 15 43 34" src="https://github.com/user-attachments/assets/b0618f96-4dea-4cab-aa31-36cc99f32194" /> 3) The pipeline runs successfully indicating that authentication succeeded. <img width="1912" height="987" alt="Screenshot 2025-09-18 at 15 43 59" src="https://github.com/user-attachments/assets/6eb5d0e9-bc04-414e-8c04-a21cf8dce078" /> --------- Signed-off-by: Divyansh Vijayvergia <[email protected]>
1 parent 64a5550 commit 88f1047

File tree

5 files changed

+393
-13
lines changed

5 files changed

+393
-13
lines changed

NEXT_CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
### New Features and Improvements
66

7+
* Add native support for authentication through Azure DevOps OIDC
8+
79
### Bug Fixes
810

911
### Documentation

README.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,10 +126,11 @@ Depending on the Databricks authentication method, the SDK uses the following in
126126

127127
### Databricks native authentication
128128

129-
By default, the Databricks SDK for Python initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Databricks Workload Identity Federation (WIF) authentication using OIDC (`auth_type="github-oidc"` argument).
129+
By default, the Databricks SDK for Python initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF). See [Supported WIF](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-provider) for the supported JWT token providers.
130130

131131
- For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents.
132132
- For Databricks OIDC authentication, you must provide the `host`, `client_id` and `token_audience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file.
133+
- For Azure DevOps OIDC authentication, the `token_audience` is irrelevant as the audience is always set to `api://AzureADTokenExchange`. Also, the `System.AccessToken` pipeline variable required for OIDC request must be exposed as the `SYSTEM_ACCESSTOKEN` environment variable, following [Pipeline variables](https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken)
133134

134135
| Argument | Description | Environment variable |
135136
|------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------|

databricks/sdk/credentials_provider.py

Lines changed: 61 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import threading
1313
import time
1414
from datetime import datetime
15-
from typing import Callable, Dict, List, Optional, Tuple, Union
15+
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
1616

1717
import google.auth # type: ignore
1818
import requests
@@ -89,7 +89,6 @@ def inner(
8989
@functools.wraps(func)
9090
def wrapper(cfg: "Config") -> Optional[CredentialsProvider]:
9191
for attr in require:
92-
getattr(cfg, attr)
9392
if not getattr(cfg, attr):
9493
return None
9594
return func(cfg)
@@ -103,7 +102,12 @@ def wrapper(cfg: "Config") -> Optional[CredentialsProvider]:
103102
def oauth_credentials_strategy(name: str, require: List[str]):
104103
"""Given the function that receives a Config and returns an OauthHeaderFactory,
105104
create an OauthCredentialsProvider with a given name and required configuration
106-
attribute names to be present for this function to be called."""
105+
attribute names to be present for this function to be called.
106+
107+
Args:
108+
name: The name of the authentication strategy
109+
require: List of config attributes that must be present
110+
"""
107111

108112
def inner(
109113
func: Callable[["Config"], OAuthCredentialsProvider],
@@ -356,33 +360,47 @@ def token() -> oauth.Token:
356360
return OAuthCredentialsProvider(refreshed_headers, token)
357361

358362

359-
@oauth_credentials_strategy("github-oidc", ["host", "client_id"])
360-
def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]:
363+
def _oidc_credentials_provider(
364+
cfg: "Config", supplier_factory: Callable[[], Any], provider_name: str
365+
) -> Optional[CredentialsProvider]:
361366
"""
362-
DatabricksWIFCredentials uses a Token Supplier to get a JWT Token and exchanges
363-
it for a Databricks Token.
367+
Generic OIDC credentials provider that works with any OIDC token supplier.
368+
369+
Args:
370+
cfg: Databricks configuration
371+
supplier_factory: Callable that returns an OIDC token supplier instance
372+
provider_name: Human-readable name (e.g., "GitHub OIDC", "Azure DevOps OIDC")
364373
365-
Supported suppliers:
366-
- GitHub OIDC
374+
Returns:
375+
OAuthCredentialsProvider if successful, None if supplier unavailable or token retrieval fails
367376
"""
368-
supplier = oidc_token_supplier.GitHubOIDCTokenSupplier()
377+
# Try to create the supplier
378+
try:
379+
supplier = supplier_factory()
380+
except Exception as e:
381+
logger.debug(f"{provider_name}: {str(e)}")
382+
return None
369383

384+
# Determine the audience for token exchange
370385
audience = cfg.token_audience
371386
if audience is None and cfg.is_account_client:
372387
audience = cfg.account_id
373388
if audience is None and not cfg.is_account_client:
374389
audience = cfg.oidc_endpoints.token_endpoint
375390

376-
# Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode.
391+
# Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode.
377392
id_token = supplier.get_oidc_token(audience)
378393
if not id_token:
394+
logger.debug(f"{provider_name}: no token available, skipping authentication method")
379395
return None
380396

397+
logger.info(f"Configured {provider_name} authentication")
398+
381399
def token_source_for(audience: str) -> oauth.TokenSource:
382400
id_token = supplier.get_oidc_token(audience)
383401
if not id_token:
384402
# Should not happen, since we checked it above.
385-
raise Exception("Cannot get OIDC token")
403+
raise Exception(f"Cannot get {provider_name} token")
386404

387405
return oauth.ClientCredentials(
388406
client_id=cfg.client_id,
@@ -408,6 +426,36 @@ def token() -> oauth.Token:
408426
return OAuthCredentialsProvider(refreshed_headers, token)
409427

410428

429+
@oauth_credentials_strategy("github-oidc", ["host", "client_id"])
430+
def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]:
431+
"""
432+
GitHub OIDC authentication uses a Token Supplier to get a JWT Token and exchanges
433+
it for a Databricks Token.
434+
435+
Supported in GitHub Actions with OIDC service connections.
436+
"""
437+
return _oidc_credentials_provider(
438+
cfg=cfg,
439+
supplier_factory=lambda: oidc_token_supplier.GitHubOIDCTokenSupplier(),
440+
provider_name="GitHub OIDC",
441+
)
442+
443+
444+
@oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"])
445+
def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]:
446+
"""
447+
Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token
448+
and exchanges it for a Databricks Token.
449+
450+
Supported in Azure DevOps pipelines with OIDC service connections.
451+
"""
452+
return _oidc_credentials_provider(
453+
cfg=cfg,
454+
supplier_factory=lambda: oidc_token_supplier.AzureDevOpsOIDCTokenSupplier(),
455+
provider_name="Azure DevOps OIDC",
456+
)
457+
458+
411459
@oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"])
412460
def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
413461
if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ:
@@ -1019,6 +1067,7 @@ def __init__(self) -> None:
10191067
azure_service_principal,
10201068
github_oidc_azure,
10211069
azure_cli,
1070+
azure_devops_oidc,
10221071
external_browser,
10231072
databricks_cli,
10241073
runtime_native_auth,

databricks/sdk/oidc_token_supplier.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
1+
import logging
12
import os
23
from typing import Optional
34

45
import requests
56

7+
logger = logging.getLogger("databricks.sdk")
68

9+
10+
# TODO: Check the required environment variables while creating the instance rather than in the get_oidc_token method to allow early return.
711
class GitHubOIDCTokenSupplier:
812
"""
913
Supplies OIDC tokens from GitHub Actions.
@@ -26,3 +30,79 @@ def get_oidc_token(self, audience: str) -> Optional[str]:
2630
return None
2731

2832
return response_json["value"]
33+
34+
35+
class AzureDevOpsOIDCTokenSupplier:
36+
"""
37+
Supplies OIDC tokens from Azure DevOps pipelines.
38+
39+
Constructs the OIDC token request URL using official Azure DevOps predefined variables.
40+
See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables
41+
"""
42+
43+
def __init__(self):
44+
"""Initialize and validate Azure DevOps environment variables."""
45+
# Get Azure DevOps environment variables.
46+
self.access_token = os.environ.get("SYSTEM_ACCESSTOKEN")
47+
self.collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI")
48+
self.project_id = os.environ.get("SYSTEM_TEAMPROJECTID")
49+
self.plan_id = os.environ.get("SYSTEM_PLANID")
50+
self.job_id = os.environ.get("SYSTEM_JOBID")
51+
self.hub_name = os.environ.get("SYSTEM_HOSTTYPE")
52+
53+
# Check for required variables with specific error messages.
54+
missing_vars = []
55+
if not self.access_token:
56+
missing_vars.append("SYSTEM_ACCESSTOKEN")
57+
if not self.collection_uri:
58+
missing_vars.append("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI")
59+
if not self.project_id:
60+
missing_vars.append("SYSTEM_TEAMPROJECTID")
61+
if not self.plan_id:
62+
missing_vars.append("SYSTEM_PLANID")
63+
if not self.job_id:
64+
missing_vars.append("SYSTEM_JOBID")
65+
if not self.hub_name:
66+
missing_vars.append("SYSTEM_HOSTTYPE")
67+
68+
if missing_vars:
69+
if "SYSTEM_ACCESSTOKEN" in missing_vars:
70+
error_msg = "Azure DevOps OIDC: SYSTEM_ACCESSTOKEN env var not found. If calling from Azure DevOps Pipeline, please set this env var following https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken"
71+
else:
72+
error_msg = f"Azure DevOps OIDC: missing required environment variables: {', '.join(missing_vars)}"
73+
raise ValueError(error_msg)
74+
75+
def get_oidc_token(self, audience: str) -> Optional[str]:
76+
# Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange".
77+
# The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers.
78+
79+
try:
80+
# Construct the OIDC token request URL.
81+
# Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken.
82+
request_url = f"{self.collection_uri}{self.project_id}/_apis/distributedtask/hubs/{self.hub_name}/plans/{self.plan_id}/jobs/{self.job_id}/oidctoken"
83+
84+
# Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps).
85+
endpoint = f"{request_url}?api-version=7.2-preview.1"
86+
headers = {
87+
"Authorization": f"Bearer {self.access_token}",
88+
"Content-Type": "application/json",
89+
"Content-Length": "0",
90+
}
91+
92+
# Azure DevOps OIDC endpoint requires POST request with empty body.
93+
response = requests.post(endpoint, headers=headers)
94+
if not response.ok:
95+
logger.debug(f"Azure DevOps OIDC: token request failed with status {response.status_code}")
96+
return None
97+
98+
# Azure DevOps returns the token in 'oidcToken' field.
99+
response_json = response.json()
100+
if "oidcToken" not in response_json:
101+
logger.debug("Azure DevOps OIDC: response missing 'oidcToken' field")
102+
return None
103+
104+
logger.debug("Azure DevOps OIDC: successfully obtained token")
105+
return response_json["oidcToken"]
106+
except Exception as e:
107+
logger.debug(f"Azure DevOps OIDC: failed to get token: {e}")
108+
return None

0 commit comments

Comments
 (0)