From c1c95cba05156fcb173a0be2307663c646bf3e00 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 14 Aug 2025 09:01:55 +0000 Subject: [PATCH 01/19] added Azure DevOps OIDC authentication --- databricks/sdk/credentials_provider.py | 52 ++++++++++++++++++++++++++ databricks/sdk/oidc_token_supplier.py | 45 ++++++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 613172cf1..564287da8 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -407,6 +407,57 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) +@oauth_credentials_strategy("azdo-oidc", ["host", "client_id"]) +def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: + """ + Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token + and exchanges it for a Databricks Token. + + Supported in Azure DevOps pipelines with OIDC service connections. + """ + supplier = oidc_token_supplier.AzureDevOpsOIDCTokenSupplier() + + audience = cfg.token_audience + if audience is None and cfg.is_account_client: + audience = cfg.account_id + if audience is None and not cfg.is_account_client: + audience = cfg.oidc_endpoints.token_endpoint + + # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. + id_token = supplier.get_oidc_token(audience) + if not id_token: + return None + + def token_source_for(audience: str) -> oauth.TokenSource: + id_token = supplier.get_oidc_token(audience) + if not id_token: + # Should not happen, since we checked it above. + raise Exception("Cannot get Azure DevOps OIDC token") + + return oauth.ClientCredentials( + client_id=cfg.client_id, + client_secret="", # we have no (rotatable) secrets in OIDC flow + token_url=cfg.oidc_endpoints.token_endpoint, + endpoint_params={ + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "subject_token": id_token, + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + }, + scopes=["all-apis"], + use_params=True, + disable_async=cfg.disable_async_token_refresh, + ) + + def refreshed_headers() -> Dict[str, str]: + token = token_source_for(audience).token() + return {"Authorization": f"{token.token_type} {token.access_token}"} + + def token() -> oauth.Token: + return token_source_for(audience).token() + + return OAuthCredentialsProvider(refreshed_headers, token) + + @oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"]) def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]: @@ -1016,6 +1067,7 @@ def __init__(self) -> None: env_oidc, file_oidc, github_oidc, + azure_devops_oidc, azure_service_principal, github_oidc_azure, azure_cli, diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index dfd139de5..9007a2519 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -26,3 +26,48 @@ def get_oidc_token(self, audience: str) -> Optional[str]: return None return response_json["value"] + + +class AzureDevOpsOIDCTokenSupplier: + """ + Supplies OIDC tokens from Azure DevOps pipelines. + """ + + def get_oidc_token(self, audience: str) -> Optional[str]: + # Retrieve necessary environment variables + access_token = os.environ.get('SYSTEM_ACCESSTOKEN') + collection_uri = os.environ.get('SYSTEM_TEAMFOUNDATIONCOLLECTIONURI') + project_id = os.environ.get('SYSTEM_TEAMPROJECT') + plan_id = os.environ.get('SYSTEM_PLANID') + job_id = os.environ.get('SYSTEM_JOBID') + hub_name = 'build' # Or 'release' for release pipelines + + # Check for required variables + if not all([access_token, collection_uri, project_id, plan_id, job_id]): + # not in Azure DevOps pipeline + return None + + # Construct the URL + request_url = f"{collection_uri}{project_id}/_apis/distributedtask/hubs/{hub_name}/plans/{plan_id}/jobs/{job_id}/oidctoken?api-version=7.1-preview.1" + + # See https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json" + } + + # Set up the JSON payload for the request body + payload = { + 'audience': audience + } + try: + # Make the POST request + response = requests.post(request_url, headers=headers, json=payload) + response.raise_for_status() # Raise an exception for bad status codes + + # Return the OIDC token from the JSON response + return response.json()['oidcToken'] + except requests.exceptions.RequestException as e: + print(f"Error requesting OIDC token: {e}") + return None + \ No newline at end of file From 8e4bb86915a4f211debc7f78dcef9faec3d9508d Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 14 Aug 2025 11:03:10 +0000 Subject: [PATCH 02/19] improved naming convention and removed redundant code --- databricks/sdk/credentials_provider.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 564287da8..bebc4cc63 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -89,7 +89,6 @@ def inner( @functools.wraps(func) def wrapper(cfg: "Config") -> Optional[CredentialsProvider]: for attr in require: - getattr(cfg, attr) if not getattr(cfg, attr): return None return func(cfg) @@ -1055,11 +1054,11 @@ def model_serving_auth(cfg: "Config") -> Optional[CredentialsProvider]: class DefaultCredentials: - """Select the first applicable credential provider from the chain""" + """Select the first applicable credential strategy from the chain""" def __init__(self) -> None: self._auth_type = "default" - self._auth_providers = [ + self._auth_strategies = [ pat_auth, basic_auth, metadata_service, @@ -1083,26 +1082,26 @@ def auth_type(self) -> str: return self._auth_type def oauth_token(self, cfg: "Config") -> oauth.Token: - for provider in self._auth_providers: - auth_type = provider.auth_type() + for strategy in self._auth_strategies: + auth_type = strategy.auth_type() if auth_type != self._auth_type: # ignore other auth types if they don't match the selected one continue - return provider.oauth_token(cfg) + return strategy.oauth_token(cfg) def __call__(self, cfg: "Config") -> CredentialsProvider: - for provider in self._auth_providers: - auth_type = provider.auth_type() + for strategy in self._auth_strategies: + auth_type = strategy.auth_type() if cfg.auth_type and auth_type != cfg.auth_type: # ignore other auth types if one is explicitly enforced logger.debug(f"Ignoring {auth_type} auth, because {cfg.auth_type} is preferred") continue logger.debug(f"Attempting to configure auth: {auth_type}") try: - # The header factory might be None if the provider cannot be + # The header factory might be None if the strategy cannot be # configured for the current environment. For example, if the - # provider requires some missing environment variables. - header_factory = provider(cfg) + # strategy requires some missing environment variables. + header_factory = strategy(cfg) if not header_factory: continue From 751c1a4b959b5ef79d4d23a1413d2545f6431aae Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 09:30:40 +0000 Subject: [PATCH 03/19] corrected azure devops OIDC token request --- databricks/sdk/oidc_token_supplier.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 9007a2519..931298c69 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -37,10 +37,10 @@ def get_oidc_token(self, audience: str) -> Optional[str]: # Retrieve necessary environment variables access_token = os.environ.get('SYSTEM_ACCESSTOKEN') collection_uri = os.environ.get('SYSTEM_TEAMFOUNDATIONCOLLECTIONURI') - project_id = os.environ.get('SYSTEM_TEAMPROJECT') + project_id = os.environ.get('SYSTEM_TEAMPROJECTID') plan_id = os.environ.get('SYSTEM_PLANID') job_id = os.environ.get('SYSTEM_JOBID') - hub_name = 'build' # Or 'release' for release pipelines + hub_name = os.environ.get('SYSTEM_HOSTTYPE') # Check for required variables if not all([access_token, collection_uri, project_id, plan_id, job_id]): From d4175ae5ad22bbe8e16c3bf463722331a2cc0d56 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 12:17:21 +0000 Subject: [PATCH 04/19] removed audience, used latest api version --- databricks/sdk/oidc_token_supplier.py | 65 +++++++++++++++------------ 1 file changed, 37 insertions(+), 28 deletions(-) diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 931298c69..a34f72626 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -31,43 +31,52 @@ def get_oidc_token(self, audience: str) -> Optional[str]: class AzureDevOpsOIDCTokenSupplier: """ Supplies OIDC tokens from Azure DevOps pipelines. + + Constructs the OIDC token request URL using official Azure DevOps predefined variables. + See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ def get_oidc_token(self, audience: str) -> Optional[str]: - # Retrieve necessary environment variables - access_token = os.environ.get('SYSTEM_ACCESSTOKEN') - collection_uri = os.environ.get('SYSTEM_TEAMFOUNDATIONCOLLECTIONURI') - project_id = os.environ.get('SYSTEM_TEAMPROJECTID') - plan_id = os.environ.get('SYSTEM_PLANID') - job_id = os.environ.get('SYSTEM_JOBID') - hub_name = os.environ.get('SYSTEM_HOSTTYPE') + # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" + # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers + + # Check for required Azure DevOps environment variables + access_token = os.environ.get("SYSTEM_ACCESSTOKEN") + collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") + project_id = os.environ.get("SYSTEM_TEAMPROJECTID") + plan_id = os.environ.get("SYSTEM_PLANID") + job_id = os.environ.get("SYSTEM_JOBID") + hub_name = os.environ.get("SYSTEM_HOSTTYPE", "build") # Default to "build" # Check for required variables if not all([access_token, collection_uri, project_id, plan_id, job_id]): # not in Azure DevOps pipeline return None - # Construct the URL - request_url = f"{collection_uri}{project_id}/_apis/distributedtask/hubs/{hub_name}/plans/{plan_id}/jobs/{job_id}/oidctoken?api-version=7.1-preview.1" - - # See https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables - headers = { - "Authorization": f"Bearer {access_token}", - "Content-Type": "application/json" - } - - # Set up the JSON payload for the request body - payload = { - 'audience': audience - } try: - # Make the POST request - response = requests.post(request_url, headers=headers, json=payload) - response.raise_for_status() # Raise an exception for bad status codes + # Construct the OIDC token request URL + # Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken + request_url = f"{collection_uri}{project_id}/_apis/distributedtask/hubs/{hub_name}/plans/{plan_id}/jobs/{job_id}/oidctoken" + + # Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps) + endpoint = f"{request_url}?api-version=7.2-preview.1" + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + "Content-Length": "0" + } + + # Azure DevOps OIDC endpoint requires POST request with empty body + response = requests.post(endpoint, headers=headers) + if not response.ok: + return None + + # Azure DevOps returns the token in 'oidcToken' field + response_json = response.json() + if "oidcToken" not in response_json: + return None - # Return the OIDC token from the JSON response - return response.json()['oidcToken'] - except requests.exceptions.RequestException as e: - print(f"Error requesting OIDC token: {e}") + return response_json["oidcToken"] + except Exception: + # If any error occurs, return None to fall back to other auth methods return None - \ No newline at end of file From 1bbc3c33e65f732596b56e5e73381315654017d1 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 12:23:42 +0000 Subject: [PATCH 05/19] formatting fixed --- databricks/sdk/oidc_token_supplier.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index a34f72626..247815be3 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -31,7 +31,7 @@ def get_oidc_token(self, audience: str) -> Optional[str]: class AzureDevOpsOIDCTokenSupplier: """ Supplies OIDC tokens from Azure DevOps pipelines. - + Constructs the OIDC token request URL using official Azure DevOps predefined variables. See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ @@ -39,7 +39,7 @@ class AzureDevOpsOIDCTokenSupplier: def get_oidc_token(self, audience: str) -> Optional[str]: # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers - + # Check for required Azure DevOps environment variables access_token = os.environ.get("SYSTEM_ACCESSTOKEN") collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") @@ -52,20 +52,20 @@ def get_oidc_token(self, audience: str) -> Optional[str]: if not all([access_token, collection_uri, project_id, plan_id, job_id]): # not in Azure DevOps pipeline return None - + try: # Construct the OIDC token request URL # Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken request_url = f"{collection_uri}{project_id}/_apis/distributedtask/hubs/{hub_name}/plans/{plan_id}/jobs/{job_id}/oidctoken" - + # Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps) endpoint = f"{request_url}?api-version=7.2-preview.1" headers = { "Authorization": f"Bearer {access_token}", "Content-Type": "application/json", - "Content-Length": "0" + "Content-Length": "0", } - + # Azure DevOps OIDC endpoint requires POST request with empty body response = requests.post(endpoint, headers=headers) if not response.ok: From fd24dc43850242cc0ee15c1343e381d3db2b2277 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 12:26:57 +0000 Subject: [PATCH 06/19] formatting fixed --- databricks/sdk/credentials_provider.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index bebc4cc63..dd2cf7894 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -150,9 +150,7 @@ def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: # This import MUST be after the "DATABRICKS_RUNTIME_VERSION" check # above, so that we are not throwing import errors when not in # runtime and no config variables are set. - from databricks.sdk.runtime import (init_runtime_legacy_auth, - init_runtime_native_auth, - init_runtime_repl_auth) + from databricks.sdk.runtime import init_runtime_legacy_auth, init_runtime_native_auth, init_runtime_repl_auth for init in [ init_runtime_native_auth, @@ -406,10 +404,11 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) + @oauth_credentials_strategy("azdo-oidc", ["host", "client_id"]) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ - Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token + Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token and exchanges it for a Databricks Token. Supported in Azure DevOps pipelines with OIDC service connections. @@ -457,7 +456,6 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) - @oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"]) def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]: if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ: From 553fd36dc8eb6999ee2f29848fbab2be545d8bc3 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 12:44:39 +0000 Subject: [PATCH 07/19] formatted using make fmt --- databricks/sdk/credentials_provider.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index dd2cf7894..14e280216 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -150,7 +150,9 @@ def runtime_native_auth(cfg: "Config") -> Optional[CredentialsProvider]: # This import MUST be after the "DATABRICKS_RUNTIME_VERSION" check # above, so that we are not throwing import errors when not in # runtime and no config variables are set. - from databricks.sdk.runtime import init_runtime_legacy_auth, init_runtime_native_auth, init_runtime_repl_auth + from databricks.sdk.runtime import (init_runtime_legacy_auth, + init_runtime_native_auth, + init_runtime_repl_auth) for init in [ init_runtime_native_auth, From 5bbb16b46ec86361d422a07175a03a31887de878 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 13:28:42 +0000 Subject: [PATCH 08/19] edded entry in changelog --- NEXT_CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 869a88d02..9a302881f 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,6 +6,8 @@ * Add a public helper function to build a `CredentialsProvider` directly from an `IdTokenSource`. +* Add native support for authentication through Azure DevOps OIDC + ### Bug Fixes ### Documentation From 649cc5db4240ece76fb6ece7dfec3668d5a0698d Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 16:22:04 +0000 Subject: [PATCH 09/19] added feature in Readme --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 58a885307..7f7d349ae 100644 --- a/README.md +++ b/README.md @@ -126,10 +126,11 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -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). +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) based authentication with support for GitHub Actions (`auth_type: "github-oidc"`) and Azure DevOps Pipelines (`auth_type: "azure-devops-oidc"`). - For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents. - 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. +- 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) | Argument | Description | Environment variable | |------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------| From ef7a0581d0cedce499d5e4b996f3d4726c21a0b3 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 18 Sep 2025 17:27:20 +0000 Subject: [PATCH 10/19] updated readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7f7d349ae..af8a960cb 100644 --- a/README.md +++ b/README.md @@ -126,7 +126,7 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -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) based authentication with support for GitHub Actions (`auth_type: "github-oidc"`) and Azure DevOps Pipelines (`auth_type: "azure-devops-oidc"`). +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. - For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents. - 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. From eabfc73756a819e7962ac77b61e60a60e3391315 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Fri, 19 Sep 2025 14:29:20 +0000 Subject: [PATCH 11/19] added environment variables to config --- databricks/sdk/config.py | 9 +++++++++ databricks/sdk/credentials_provider.py | 17 ++++++++++++++--- databricks/sdk/oidc_token_supplier.py | 19 +++++++++++-------- 3 files changed, 34 insertions(+), 11 deletions(-) diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 4cfd8b4f9..630a1110a 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -88,6 +88,15 @@ class Config: azure_client_id: str = ConfigAttribute(env="ARM_CLIENT_ID", auth="azure") azure_tenant_id: str = ConfigAttribute(env="ARM_TENANT_ID", auth="azure") azure_environment: str = ConfigAttribute(env="ARM_ENVIRONMENT") + + # Azure DevOps environment variables (automatically provided by Azure DevOps pipelines) + azure_devops_access_token: str = ConfigAttribute(env="SYSTEM_ACCESSTOKEN", auth="azdo-oidc", sensitive=True) + azure_devops_collection_uri: str = ConfigAttribute(env="SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", auth="azdo-oidc") + azure_devops_project_id: str = ConfigAttribute(env="SYSTEM_TEAMPROJECTID", auth="azdo-oidc") + azure_devops_plan_id: str = ConfigAttribute(env="SYSTEM_PLANID", auth="azdo-oidc") + azure_devops_job_id: str = ConfigAttribute(env="SYSTEM_JOBID", auth="azdo-oidc") + azure_devops_host_type: str = ConfigAttribute(env="SYSTEM_HOSTTYPE", auth="azdo-oidc") + databricks_cli_path: str = ConfigAttribute(env="DATABRICKS_CLI_PATH") auth_type: str = ConfigAttribute(env="DATABRICKS_AUTH_TYPE") cluster_id: str = ConfigAttribute(env="DATABRICKS_CLUSTER_ID") diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 14e280216..a6c7a5675 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -407,7 +407,18 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) -@oauth_credentials_strategy("azdo-oidc", ["host", "client_id"]) +@oauth_credentials_strategy( + "azdo-oidc", + [ + "host", + "client_id", + "azure_devops_access_token", + "azure_devops_collection_uri", + "azure_devops_project_id", + "azure_devops_plan_id", + "azure_devops_job_id", + ], +) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token @@ -424,12 +435,12 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: audience = cfg.oidc_endpoints.token_endpoint # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. - id_token = supplier.get_oidc_token(audience) + id_token = supplier.get_oidc_token(audience, cfg) if not id_token: return None def token_source_for(audience: str) -> oauth.TokenSource: - id_token = supplier.get_oidc_token(audience) + id_token = supplier.get_oidc_token(audience, cfg) if not id_token: # Should not happen, since we checked it above. raise Exception("Cannot get Azure DevOps OIDC token") diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 247815be3..209240e4f 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -36,17 +36,20 @@ class AzureDevOpsOIDCTokenSupplier: See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ - def get_oidc_token(self, audience: str) -> Optional[str]: + def get_oidc_token(self, audience: str, config=None) -> Optional[str]: # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers - # Check for required Azure DevOps environment variables - access_token = os.environ.get("SYSTEM_ACCESSTOKEN") - collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") - project_id = os.environ.get("SYSTEM_TEAMPROJECTID") - plan_id = os.environ.get("SYSTEM_PLANID") - job_id = os.environ.get("SYSTEM_JOBID") - hub_name = os.environ.get("SYSTEM_HOSTTYPE", "build") # Default to "build" + # Get Azure DevOps environment variables from config + if config is None: + return None + + access_token = config.azure_devops_access_token + collection_uri = config.azure_devops_collection_uri + project_id = config.azure_devops_project_id + plan_id = config.azure_devops_plan_id + job_id = config.azure_devops_job_id + hub_name = config.azure_devops_host_type or "build" # Default to "build" # Check for required variables if not all([access_token, collection_uri, project_id, plan_id, job_id]): From 8c95cd95c244ab399524ad50d7a8e6ebe8bcf0ff Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Fri, 19 Sep 2025 14:31:52 +0000 Subject: [PATCH 12/19] changed strategy to provider --- databricks/sdk/credentials_provider.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index a6c7a5675..0143320a7 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -1065,11 +1065,11 @@ def model_serving_auth(cfg: "Config") -> Optional[CredentialsProvider]: class DefaultCredentials: - """Select the first applicable credential strategy from the chain""" + """Select the first applicable credential provider from the chain""" def __init__(self) -> None: self._auth_type = "default" - self._auth_strategies = [ + self._auth_providers = [ pat_auth, basic_auth, metadata_service, @@ -1093,26 +1093,26 @@ def auth_type(self) -> str: return self._auth_type def oauth_token(self, cfg: "Config") -> oauth.Token: - for strategy in self._auth_strategies: - auth_type = strategy.auth_type() + for provider in self._auth_providers: + auth_type = provider.auth_type() if auth_type != self._auth_type: # ignore other auth types if they don't match the selected one continue - return strategy.oauth_token(cfg) + return provider.oauth_token(cfg) def __call__(self, cfg: "Config") -> CredentialsProvider: - for strategy in self._auth_strategies: - auth_type = strategy.auth_type() + for provider in self._auth_providers: + auth_type = provider.auth_type() if cfg.auth_type and auth_type != cfg.auth_type: # ignore other auth types if one is explicitly enforced logger.debug(f"Ignoring {auth_type} auth, because {cfg.auth_type} is preferred") continue logger.debug(f"Attempting to configure auth: {auth_type}") try: - # The header factory might be None if the strategy cannot be + # The header factory might be None if the provider cannot be # configured for the current environment. For example, if the - # strategy requires some missing environment variables. - header_factory = strategy(cfg) + # provider requires some missing environment variables. + header_factory = provider(cfg) if not header_factory: continue From 44c78c16e491f97a71d52cb1990c7b62011055ad Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Tue, 23 Sep 2025 14:44:07 +0000 Subject: [PATCH 13/19] added tests, removed env vars from config, early fail --- databricks/sdk/config.py | 9 -- databricks/sdk/credentials_provider.py | 31 +++-- databricks/sdk/oidc_token_supplier.py | 23 ++-- tests/test_oidc_token_supplier.py | 172 +++++++++++++++++++++++++ 4 files changed, 199 insertions(+), 36 deletions(-) create mode 100644 tests/test_oidc_token_supplier.py diff --git a/databricks/sdk/config.py b/databricks/sdk/config.py index 630a1110a..4cfd8b4f9 100644 --- a/databricks/sdk/config.py +++ b/databricks/sdk/config.py @@ -88,15 +88,6 @@ class Config: azure_client_id: str = ConfigAttribute(env="ARM_CLIENT_ID", auth="azure") azure_tenant_id: str = ConfigAttribute(env="ARM_TENANT_ID", auth="azure") azure_environment: str = ConfigAttribute(env="ARM_ENVIRONMENT") - - # Azure DevOps environment variables (automatically provided by Azure DevOps pipelines) - azure_devops_access_token: str = ConfigAttribute(env="SYSTEM_ACCESSTOKEN", auth="azdo-oidc", sensitive=True) - azure_devops_collection_uri: str = ConfigAttribute(env="SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", auth="azdo-oidc") - azure_devops_project_id: str = ConfigAttribute(env="SYSTEM_TEAMPROJECTID", auth="azdo-oidc") - azure_devops_plan_id: str = ConfigAttribute(env="SYSTEM_PLANID", auth="azdo-oidc") - azure_devops_job_id: str = ConfigAttribute(env="SYSTEM_JOBID", auth="azdo-oidc") - azure_devops_host_type: str = ConfigAttribute(env="SYSTEM_HOSTTYPE", auth="azdo-oidc") - databricks_cli_path: str = ConfigAttribute(env="DATABRICKS_CLI_PATH") auth_type: str = ConfigAttribute(env="DATABRICKS_AUTH_TYPE") cluster_id: str = ConfigAttribute(env="DATABRICKS_CLUSTER_ID") diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 0143320a7..dee60869f 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -99,16 +99,26 @@ def wrapper(cfg: "Config") -> Optional[CredentialsProvider]: return inner -def oauth_credentials_strategy(name: str, require: List[str]): +def oauth_credentials_strategy(name: str, require: List[str], env_vars: Optional[List[str]] = None): """Given the function that receives a Config and returns an OauthHeaderFactory, create an OauthCredentialsProvider with a given name and required configuration - attribute names to be present for this function to be called.""" + attribute names to be present for this function to be called. + + Args: + name: The name of the authentication strategy + require: List of config attributes that must be present + env_vars: Optional list of environment variables that must all be present for this strategy + """ def inner( func: Callable[["Config"], OAuthCredentialsProvider], ) -> OauthCredentialsStrategy: @functools.wraps(func) def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: + # Early environment detection - check before config validation + if env_vars and not all(os.environ.get(var) for var in env_vars): + return None + for attr in require: if not getattr(cfg, attr): return None @@ -408,16 +418,9 @@ def token() -> oauth.Token: @oauth_credentials_strategy( - "azdo-oidc", - [ - "host", - "client_id", - "azure_devops_access_token", - "azure_devops_collection_uri", - "azure_devops_project_id", - "azure_devops_plan_id", - "azure_devops_job_id", - ], + "azdo-oidc", + ["host", "client_id"], + env_vars=["SYSTEM_ACCESSTOKEN", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", "SYSTEM_TEAMPROJECTID", "SYSTEM_PLANID", "SYSTEM_JOBID", "SYSTEM_HOSTTYPE"] ) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ @@ -435,12 +438,12 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: audience = cfg.oidc_endpoints.token_endpoint # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. - id_token = supplier.get_oidc_token(audience, cfg) + id_token = supplier.get_oidc_token(audience) if not id_token: return None def token_source_for(audience: str) -> oauth.TokenSource: - id_token = supplier.get_oidc_token(audience, cfg) + id_token = supplier.get_oidc_token(audience) if not id_token: # Should not happen, since we checked it above. raise Exception("Cannot get Azure DevOps OIDC token") diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 209240e4f..6ac58fa5e 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -31,28 +31,25 @@ def get_oidc_token(self, audience: str) -> Optional[str]: class AzureDevOpsOIDCTokenSupplier: """ Supplies OIDC tokens from Azure DevOps pipelines. - + Constructs the OIDC token request URL using official Azure DevOps predefined variables. See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ - def get_oidc_token(self, audience: str, config=None) -> Optional[str]: + def get_oidc_token(self, audience: str) -> Optional[str]: # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers + - # Get Azure DevOps environment variables from config - if config is None: - return None - - access_token = config.azure_devops_access_token - collection_uri = config.azure_devops_collection_uri - project_id = config.azure_devops_project_id - plan_id = config.azure_devops_plan_id - job_id = config.azure_devops_job_id - hub_name = config.azure_devops_host_type or "build" # Default to "build" + access_token = os.environ.get("SYSTEM_ACCESSTOKEN") + collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") + project_id = os.environ.get("SYSTEM_TEAMPROJECTID") + plan_id = os.environ.get("SYSTEM_PLANID") + job_id = os.environ.get("SYSTEM_JOBID") + hub_name = os.environ.get("SYSTEM_HOSTTYPE") # Check for required variables - if not all([access_token, collection_uri, project_id, plan_id, job_id]): + if not all([access_token, collection_uri, project_id, plan_id, job_id, hub_name]): # not in Azure DevOps pipeline return None diff --git a/tests/test_oidc_token_supplier.py b/tests/test_oidc_token_supplier.py new file mode 100644 index 000000000..63aeeee59 --- /dev/null +++ b/tests/test_oidc_token_supplier.py @@ -0,0 +1,172 @@ +from dataclasses import dataclass +from typing import Dict, Optional + +import pytest + +from databricks.sdk.oidc_token_supplier import AzureDevOpsOIDCTokenSupplier + + +@dataclass +class AzureDevOpsOIDCTestCase: + name: str + env_vars: Optional[Dict[str, str]] = None + response_ok: bool = True + response_json: Optional[Dict[str, str]] = None + want_token: Optional[str] = None + want_none: bool = False + + +_azure_devops_oidc_test_cases = [ + AzureDevOpsOIDCTestCase( + name="success_with_hosttype", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + response_ok=True, + response_json={"oidcToken": "test-azdo-jwt-token"}, + want_token="test-azdo-jwt-token", + ), + AzureDevOpsOIDCTestCase( + name="missing_hosttype", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_access_token", + env_vars={ + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_plan_id", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_job_id", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_HOSTTYPE": "build" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_team_foundation_collection_uri", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_project_id", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="request_failure", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + response_ok=False, + want_none=True, + ), + AzureDevOpsOIDCTestCase( + name="missing_oidc_token_in_response", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build" + }, + response_ok=True, + response_json={"error": "no oidcToken"}, + want_none=True, + ), +] + + +@pytest.mark.parametrize("test_case", _azure_devops_oidc_test_cases) +def test_azure_devops_oidc_token_supplier(test_case: AzureDevOpsOIDCTestCase, monkeypatch, mocker): + """Test AzureDevOpsOIDCTokenSupplier with various scenarios""" + # Set up environment variables + if test_case.env_vars: + for key, value in test_case.env_vars.items(): + monkeypatch.setenv(key, value) + + # Mock requests.post if we have all required environment variables (including HOSTTYPE) + required_vars = ["SYSTEM_ACCESSTOKEN", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", + "SYSTEM_TEAMPROJECTID", "SYSTEM_PLANID", "SYSTEM_JOBID", "SYSTEM_HOSTTYPE"] + has_required_vars = test_case.env_vars and all(var in test_case.env_vars for var in required_vars) + + mock_post = None + if has_required_vars: # Only mock if all required vars exist + mock_response = mocker.Mock() + mock_response.ok = test_case.response_ok + if test_case.response_json: + mock_response.json.return_value = test_case.response_json + mock_post = mocker.patch('requests.post', return_value=mock_response) + + supplier = AzureDevOpsOIDCTokenSupplier() + token = supplier.get_oidc_token("ignored-audience") # Audience is ignored for Azure DevOps + + if test_case.want_none: + assert token is None + else: + assert token == test_case.want_token + # Verify the request was made correctly + if mock_post is not None: + expected_url = ( + "https://dev.azure.com/myorg/project-123/_apis/distributedtask/" + "hubs/build/plans/plan-456/jobs/job-789/oidctoken?api-version=7.2-preview.1" + ) + mock_post.assert_called_once_with( + expected_url, + headers={ + "Authorization": "Bearer azdo-access-token", + "Content-Type": "application/json", + "Content-Length": "0" + } + ) + From 3f6505aee6903000d1e41b90676e14b2dcba53bd Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Tue, 23 Sep 2025 14:49:55 +0000 Subject: [PATCH 14/19] ran make fmt --- databricks/sdk/credentials_provider.py | 17 ++++++--- databricks/sdk/oidc_token_supplier.py | 3 +- tests/test_oidc_token_supplier.py | 49 ++++++++++++++------------ 3 files changed, 40 insertions(+), 29 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index dee60869f..4af620f7a 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -103,7 +103,7 @@ def oauth_credentials_strategy(name: str, require: List[str], env_vars: Optional """Given the function that receives a Config and returns an OauthHeaderFactory, create an OauthCredentialsProvider with a given name and required configuration attribute names to be present for this function to be called. - + Args: name: The name of the authentication strategy require: List of config attributes that must be present @@ -118,7 +118,7 @@ def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: # Early environment detection - check before config validation if env_vars and not all(os.environ.get(var) for var in env_vars): return None - + for attr in require: if not getattr(cfg, attr): return None @@ -418,9 +418,16 @@ def token() -> oauth.Token: @oauth_credentials_strategy( - "azdo-oidc", - ["host", "client_id"], - env_vars=["SYSTEM_ACCESSTOKEN", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", "SYSTEM_TEAMPROJECTID", "SYSTEM_PLANID", "SYSTEM_JOBID", "SYSTEM_HOSTTYPE"] + "azdo-oidc", + ["host", "client_id"], + env_vars=[ + "SYSTEM_ACCESSTOKEN", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", + "SYSTEM_TEAMPROJECTID", + "SYSTEM_PLANID", + "SYSTEM_JOBID", + "SYSTEM_HOSTTYPE", + ], ) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 6ac58fa5e..f55d53884 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -31,7 +31,7 @@ def get_oidc_token(self, audience: str) -> Optional[str]: class AzureDevOpsOIDCTokenSupplier: """ Supplies OIDC tokens from Azure DevOps pipelines. - + Constructs the OIDC token request URL using official Azure DevOps predefined variables. See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ @@ -39,7 +39,6 @@ class AzureDevOpsOIDCTokenSupplier: def get_oidc_token(self, audience: str) -> Optional[str]: # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers - access_token = os.environ.get("SYSTEM_ACCESSTOKEN") collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") diff --git a/tests/test_oidc_token_supplier.py b/tests/test_oidc_token_supplier.py index 63aeeee59..a5411e886 100644 --- a/tests/test_oidc_token_supplier.py +++ b/tests/test_oidc_token_supplier.py @@ -25,7 +25,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, response_ok=True, response_json={"oidcToken": "test-azdo-jwt-token"}, @@ -36,9 +36,9 @@ class AzureDevOpsOIDCTestCase: env_vars={ "SYSTEM_ACCESSTOKEN": "azdo-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", - "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", - "SYSTEM_JOBID": "job-789" + "SYSTEM_JOBID": "job-789", }, want_none=True, ), @@ -47,9 +47,9 @@ class AzureDevOpsOIDCTestCase: env_vars={ "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", - "SYSTEM_PLANID": "plan-456", + "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, want_none=True, ), @@ -60,7 +60,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, want_none=True, ), @@ -71,7 +71,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, want_none=True, ), @@ -82,7 +82,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, want_none=True, ), @@ -93,7 +93,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, want_none=True, ), @@ -105,7 +105,7 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, response_ok=False, want_none=True, @@ -116,9 +116,9 @@ class AzureDevOpsOIDCTestCase: "SYSTEM_ACCESSTOKEN": "azdo-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", - "SYSTEM_PLANID": "plan-456", + "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build" + "SYSTEM_HOSTTYPE": "build", }, response_ok=True, response_json={"error": "no oidcToken"}, @@ -134,23 +134,29 @@ def test_azure_devops_oidc_token_supplier(test_case: AzureDevOpsOIDCTestCase, mo if test_case.env_vars: for key, value in test_case.env_vars.items(): monkeypatch.setenv(key, value) - + # Mock requests.post if we have all required environment variables (including HOSTTYPE) - required_vars = ["SYSTEM_ACCESSTOKEN", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", - "SYSTEM_TEAMPROJECTID", "SYSTEM_PLANID", "SYSTEM_JOBID", "SYSTEM_HOSTTYPE"] + required_vars = [ + "SYSTEM_ACCESSTOKEN", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", + "SYSTEM_TEAMPROJECTID", + "SYSTEM_PLANID", + "SYSTEM_JOBID", + "SYSTEM_HOSTTYPE", + ] has_required_vars = test_case.env_vars and all(var in test_case.env_vars for var in required_vars) - + mock_post = None if has_required_vars: # Only mock if all required vars exist mock_response = mocker.Mock() mock_response.ok = test_case.response_ok if test_case.response_json: mock_response.json.return_value = test_case.response_json - mock_post = mocker.patch('requests.post', return_value=mock_response) - + mock_post = mocker.patch("requests.post", return_value=mock_response) + supplier = AzureDevOpsOIDCTokenSupplier() token = supplier.get_oidc_token("ignored-audience") # Audience is ignored for Azure DevOps - + if test_case.want_none: assert token is None else: @@ -166,7 +172,6 @@ def test_azure_devops_oidc_token_supplier(test_case: AzureDevOpsOIDCTestCase, mo headers={ "Authorization": "Bearer azdo-access-token", "Content-Type": "application/json", - "Content-Length": "0" - } + "Content-Length": "0", + }, ) - From 14aa77067f514c7d1e8d154e396541f4b751f534 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Wed, 24 Sep 2025 13:37:51 +0000 Subject: [PATCH 15/19] added logging --- databricks/sdk/credentials_provider.py | 8 ++++++++ databricks/sdk/oidc_token_supplier.py | 11 +++++++++-- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 4af620f7a..82011a942 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -117,6 +117,11 @@ def inner( def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: # Early environment detection - check before config validation if env_vars and not all(os.environ.get(var) for var in env_vars): + # Provide specific error message for Azure DevOps OIDC SYSTEM_ACCESSTOKEN + if name == "azdo-oidc" and "SYSTEM_ACCESSTOKEN" in env_vars and not os.environ.get("SYSTEM_ACCESSTOKEN"): + logger.debug("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") + else: + logger.debug(f"{name}: required environment variables not present, skipping") return None for attr in require: @@ -447,8 +452,11 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. id_token = supplier.get_oidc_token(audience) if not id_token: + logger.debug("Azure DevOps OIDC: no token available, skipping authentication method") return None + logger.info("Configured Azure DevOps OIDC authentication") + def token_source_for(audience: str) -> oauth.TokenSource: id_token = supplier.get_oidc_token(audience) if not id_token: diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index f55d53884..ab8380c59 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -1,8 +1,11 @@ +import logging import os from typing import Optional import requests +logger = logging.getLogger("databricks.sdk") + class GitHubOIDCTokenSupplier: """ @@ -50,6 +53,7 @@ def get_oidc_token(self, audience: str) -> Optional[str]: # Check for required variables if not all([access_token, collection_uri, project_id, plan_id, job_id, hub_name]): # not in Azure DevOps pipeline + logger.debug("Azure DevOps OIDC: not in Azure DevOps pipeline environment") return None try: @@ -68,14 +72,17 @@ def get_oidc_token(self, audience: str) -> Optional[str]: # Azure DevOps OIDC endpoint requires POST request with empty body response = requests.post(endpoint, headers=headers) if not response.ok: + logger.debug(f"Azure DevOps OIDC: token request failed with status {response.status_code}") return None # Azure DevOps returns the token in 'oidcToken' field response_json = response.json() if "oidcToken" not in response_json: + logger.debug("Azure DevOps OIDC: response missing 'oidcToken' field") return None + logger.debug("Azure DevOps OIDC: successfully obtained token") return response_json["oidcToken"] - except Exception: - # If any error occurs, return None to fall back to other auth methods + except Exception as e: + logger.debug(f"Azure DevOps OIDC: failed to get token: {e}") return None From 3b5ab96e285ea9ac1c27e52f4a921d0fc93a20fb Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Wed, 24 Sep 2025 13:38:15 +0000 Subject: [PATCH 16/19] formatted --- databricks/sdk/credentials_provider.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 82011a942..d41d66438 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -118,8 +118,14 @@ def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: # Early environment detection - check before config validation if env_vars and not all(os.environ.get(var) for var in env_vars): # Provide specific error message for Azure DevOps OIDC SYSTEM_ACCESSTOKEN - if name == "azdo-oidc" and "SYSTEM_ACCESSTOKEN" in env_vars and not os.environ.get("SYSTEM_ACCESSTOKEN"): - logger.debug("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") + if ( + name == "azdo-oidc" + and "SYSTEM_ACCESSTOKEN" in env_vars + and not os.environ.get("SYSTEM_ACCESSTOKEN") + ): + logger.debug( + "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" + ) else: logger.debug(f"{name}: required environment variables not present, skipping") return None From 3393d78aa09eac19d1f641d6cd7bc632a48e6454 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Thu, 25 Sep 2025 18:09:28 +0000 Subject: [PATCH 17/19] restructured code and tests --- databricks/sdk/credentials_provider.py | 37 +---- databricks/sdk/oidc_token_supplier.py | 63 +++++--- tests/test_oidc_token_supplier.py | 209 +++++++++++++++++-------- 3 files changed, 188 insertions(+), 121 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index d41d66438..5b43bd389 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -99,7 +99,7 @@ def wrapper(cfg: "Config") -> Optional[CredentialsProvider]: return inner -def oauth_credentials_strategy(name: str, require: List[str], env_vars: Optional[List[str]] = None): +def oauth_credentials_strategy(name: str, require: List[str]): """Given the function that receives a Config and returns an OauthHeaderFactory, create an OauthCredentialsProvider with a given name and required configuration attribute names to be present for this function to be called. @@ -107,7 +107,6 @@ def oauth_credentials_strategy(name: str, require: List[str], env_vars: Optional Args: name: The name of the authentication strategy require: List of config attributes that must be present - env_vars: Optional list of environment variables that must all be present for this strategy """ def inner( @@ -115,21 +114,6 @@ def inner( ) -> OauthCredentialsStrategy: @functools.wraps(func) def wrapper(cfg: "Config") -> Optional[OAuthCredentialsProvider]: - # Early environment detection - check before config validation - if env_vars and not all(os.environ.get(var) for var in env_vars): - # Provide specific error message for Azure DevOps OIDC SYSTEM_ACCESSTOKEN - if ( - name == "azdo-oidc" - and "SYSTEM_ACCESSTOKEN" in env_vars - and not os.environ.get("SYSTEM_ACCESSTOKEN") - ): - logger.debug( - "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" - ) - else: - logger.debug(f"{name}: required environment variables not present, skipping") - return None - for attr in require: if not getattr(cfg, attr): return None @@ -428,18 +412,7 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) -@oauth_credentials_strategy( - "azdo-oidc", - ["host", "client_id"], - env_vars=[ - "SYSTEM_ACCESSTOKEN", - "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", - "SYSTEM_TEAMPROJECTID", - "SYSTEM_PLANID", - "SYSTEM_JOBID", - "SYSTEM_HOSTTYPE", - ], -) +@oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"]) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ Azure DevOps OIDC authentication uses a Token Supplier to get a JWT Token @@ -447,7 +420,11 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: Supported in Azure DevOps pipelines with OIDC service connections. """ - supplier = oidc_token_supplier.AzureDevOpsOIDCTokenSupplier() + try: + supplier = oidc_token_supplier.AzureDevOpsOIDCTokenSupplier() + except ValueError as e: + logger.debug(str(e)) + return None audience = cfg.token_audience if audience is None and cfg.is_account_client: diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index ab8380c59..70efd9d59 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -39,43 +39,62 @@ class AzureDevOpsOIDCTokenSupplier: See: https://docs.microsoft.com/en-us/azure/devops/pipelines/build/variables """ + def __init__(self): + """Initialize and validate Azure DevOps environment variables.""" + # Get Azure DevOps environment variables. + self.access_token = os.environ.get("SYSTEM_ACCESSTOKEN") + self.collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") + self.project_id = os.environ.get("SYSTEM_TEAMPROJECTID") + self.plan_id = os.environ.get("SYSTEM_PLANID") + self.job_id = os.environ.get("SYSTEM_JOBID") + self.hub_name = os.environ.get("SYSTEM_HOSTTYPE") + + # Check for required variables with specific error messages. + missing_vars = [] + if not self.access_token: + missing_vars.append("SYSTEM_ACCESSTOKEN") + if not self.collection_uri: + missing_vars.append("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") + if not self.project_id: + missing_vars.append("SYSTEM_TEAMPROJECTID") + if not self.plan_id: + missing_vars.append("SYSTEM_PLANID") + if not self.job_id: + missing_vars.append("SYSTEM_JOBID") + if not self.hub_name: + missing_vars.append("SYSTEM_HOSTTYPE") + + if missing_vars: + if "SYSTEM_ACCESSTOKEN" in missing_vars: + 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" + else: + error_msg = f"Azure DevOps OIDC: missing required environment variables: {', '.join(missing_vars)}" + raise ValueError(error_msg) + def get_oidc_token(self, audience: str) -> Optional[str]: - # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange" - # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers - - access_token = os.environ.get("SYSTEM_ACCESSTOKEN") - collection_uri = os.environ.get("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI") - project_id = os.environ.get("SYSTEM_TEAMPROJECTID") - plan_id = os.environ.get("SYSTEM_PLANID") - job_id = os.environ.get("SYSTEM_JOBID") - hub_name = os.environ.get("SYSTEM_HOSTTYPE") - - # Check for required variables - if not all([access_token, collection_uri, project_id, plan_id, job_id, hub_name]): - # not in Azure DevOps pipeline - logger.debug("Azure DevOps OIDC: not in Azure DevOps pipeline environment") - return None + # Note: Azure DevOps OIDC tokens have a fixed audience of "api://AzureADTokenExchange". + # The audience parameter is ignored but kept for interface compatibility with other OIDC suppliers. try: - # Construct the OIDC token request URL - # Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken - request_url = f"{collection_uri}{project_id}/_apis/distributedtask/hubs/{hub_name}/plans/{plan_id}/jobs/{job_id}/oidctoken" + # Construct the OIDC token request URL. + # Format: {collection_uri}{project_id}/_apis/distributedtask/hubs/{hubName}/plans/{planId}/jobs/{jobId}/oidctoken. + request_url = f"{self.collection_uri}{self.project_id}/_apis/distributedtask/hubs/{self.hub_name}/plans/{self.plan_id}/jobs/{self.job_id}/oidctoken" - # Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps) + # Add API version (audience is fixed to "api://AzureADTokenExchange" by Azure DevOps). endpoint = f"{request_url}?api-version=7.2-preview.1" headers = { - "Authorization": f"Bearer {access_token}", + "Authorization": f"Bearer {self.access_token}", "Content-Type": "application/json", "Content-Length": "0", } - # Azure DevOps OIDC endpoint requires POST request with empty body + # Azure DevOps OIDC endpoint requires POST request with empty body. response = requests.post(endpoint, headers=headers) if not response.ok: logger.debug(f"Azure DevOps OIDC: token request failed with status {response.status_code}") return None - # Azure DevOps returns the token in 'oidcToken' field + # Azure DevOps returns the token in 'oidcToken' field. response_json = response.json() if "oidcToken" not in response_json: logger.debug("Azure DevOps OIDC: response missing 'oidcToken' field") diff --git a/tests/test_oidc_token_supplier.py b/tests/test_oidc_token_supplier.py index a5411e886..57109c37b 100644 --- a/tests/test_oidc_token_supplier.py +++ b/tests/test_oidc_token_supplier.py @@ -7,100 +7,145 @@ @dataclass -class AzureDevOpsOIDCTestCase: +class AzureDevOpsOIDCConstructorTestCase: + """Test case for AzureDevOpsOIDCTokenSupplier constructor validation.""" + name: str env_vars: Optional[Dict[str, str]] = None + should_raise_exception: bool = False + expected_exception_message: Optional[str] = None + + +@dataclass +class AzureDevOpsOIDCTokenRequestTestCase: + """Test case for OIDC token request/response handling (assumes constructor succeeds).""" + + name: str + env_vars: Dict[str, str] # Token request tests always have all required environment variables. response_ok: bool = True response_json: Optional[Dict[str, str]] = None want_token: Optional[str] = None want_none: bool = False -_azure_devops_oidc_test_cases = [ - AzureDevOpsOIDCTestCase( - name="success_with_hosttype", +# Test cases for constructor validation (both success and failure). +_azure_devops_oidc_constructor_test_cases = [ + # Constructor success cases. + AzureDevOpsOIDCConstructorTestCase( + name="constructor_success_all_env_vars", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", "SYSTEM_HOSTTYPE": "build", }, - response_ok=True, - response_json={"oidcToken": "test-azdo-jwt-token"}, - want_token="test-azdo-jwt-token", + should_raise_exception=False, ), - AzureDevOpsOIDCTestCase( - name="missing_hosttype", + # Constructor failure cases. + AzureDevOpsOIDCConstructorTestCase( + name="missing_access_token", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: SYSTEM_ACCESSTOKEN env var not found", ), - AzureDevOpsOIDCTestCase( - name="missing_access_token", + AzureDevOpsOIDCConstructorTestCase( + name="missing_hosttype", env_vars={ + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", - "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables: SYSTEM_HOSTTYPE", ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCConstructorTestCase( name="missing_plan_id", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_JOBID": "job-789", "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables: SYSTEM_PLANID", ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCConstructorTestCase( name="missing_job_id", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables: SYSTEM_JOBID", ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCConstructorTestCase( name="missing_team_foundation_collection_uri", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables: SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCConstructorTestCase( name="missing_project_id", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_PLANID": "plan-456", "SYSTEM_JOBID": "job-789", "SYSTEM_HOSTTYPE": "build", }, - want_none=True, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables: SYSTEM_TEAMPROJECTID", ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCConstructorTestCase( + name="missing_multiple_vars", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", + }, + should_raise_exception=True, + expected_exception_message="Azure DevOps OIDC: missing required environment variables:", + ), +] + +# Test cases for OIDC token request/response handling. +_azure_devops_oidc_token_request_test_cases = [ + AzureDevOpsOIDCTokenRequestTestCase( + name="success_with_hosttype", + env_vars={ + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", + "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", + "SYSTEM_TEAMPROJECTID": "project-123", + "SYSTEM_PLANID": "plan-456", + "SYSTEM_JOBID": "job-789", + "SYSTEM_HOSTTYPE": "build", + }, + response_ok=True, + response_json={"oidcToken": "test-azure-devops-jwt-token"}, + want_token="test-azure-devops-jwt-token", + ), + AzureDevOpsOIDCTokenRequestTestCase( name="request_failure", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", @@ -110,10 +155,10 @@ class AzureDevOpsOIDCTestCase: response_ok=False, want_none=True, ), - AzureDevOpsOIDCTestCase( + AzureDevOpsOIDCTokenRequestTestCase( name="missing_oidc_token_in_response", env_vars={ - "SYSTEM_ACCESSTOKEN": "azdo-access-token", + "SYSTEM_ACCESSTOKEN": "azure-devops-access-token", "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI": "https://dev.azure.com/myorg/", "SYSTEM_TEAMPROJECTID": "project-123", "SYSTEM_PLANID": "plan-456", @@ -127,51 +172,77 @@ class AzureDevOpsOIDCTestCase: ] -@pytest.mark.parametrize("test_case", _azure_devops_oidc_test_cases) -def test_azure_devops_oidc_token_supplier(test_case: AzureDevOpsOIDCTestCase, monkeypatch, mocker): - """Test AzureDevOpsOIDCTokenSupplier with various scenarios""" - # Set up environment variables +@pytest.mark.parametrize("test_case", _azure_devops_oidc_constructor_test_cases) +def test_azure_devops_oidc_constructor_validation(test_case: AzureDevOpsOIDCConstructorTestCase, monkeypatch): + """Test AzureDevOpsOIDCTokenSupplier constructor validation with various environment variable scenarios.""" + # Set up environment variables. if test_case.env_vars: for key, value in test_case.env_vars.items(): monkeypatch.setenv(key, value) - # Mock requests.post if we have all required environment variables (including HOSTTYPE) - required_vars = [ - "SYSTEM_ACCESSTOKEN", - "SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", - "SYSTEM_TEAMPROJECTID", - "SYSTEM_PLANID", - "SYSTEM_JOBID", - "SYSTEM_HOSTTYPE", - ] - has_required_vars = test_case.env_vars and all(var in test_case.env_vars for var in required_vars) - - mock_post = None - if has_required_vars: # Only mock if all required vars exist - mock_response = mocker.Mock() - mock_response.ok = test_case.response_ok - if test_case.response_json: - mock_response.json.return_value = test_case.response_json - mock_post = mocker.patch("requests.post", return_value=mock_response) + if test_case.should_raise_exception: + # Test that constructor raises ValueError with expected message. + with pytest.raises(ValueError) as exc_info: + AzureDevOpsOIDCTokenSupplier() + + # Verify the exception message contains the expected text. + if test_case.expected_exception_message: + assert test_case.expected_exception_message in str( + exc_info.value + ), f"Exception message should contain '{test_case.expected_exception_message}', but got: {str(exc_info.value)}" + else: + # Test that constructor succeeds. + supplier = AzureDevOpsOIDCTokenSupplier() + assert supplier is not None + # Verify that all required attributes are set. + assert supplier.access_token is not None + assert supplier.collection_uri is not None + assert supplier.project_id is not None + assert supplier.plan_id is not None + assert supplier.job_id is not None + assert supplier.hub_name is not None + +@pytest.mark.parametrize("test_case", _azure_devops_oidc_token_request_test_cases) +def test_azure_devops_oidc_token_request(test_case: AzureDevOpsOIDCTokenRequestTestCase, monkeypatch, mocker): + """Test OIDC token request/response handling (assumes constructor succeeds).""" + # Set up environment variables. + for key, value in test_case.env_vars.items(): + monkeypatch.setenv(key, value) + + # Mock HTTP response. + mock_response = mocker.Mock() + mock_response.ok = test_case.response_ok + if test_case.response_json: + mock_response.json.return_value = test_case.response_json + mock_post = mocker.patch("requests.post", return_value=mock_response) + + # Initialize supplier (should succeed for these test cases since they have all required environment variables). supplier = AzureDevOpsOIDCTokenSupplier() - token = supplier.get_oidc_token("ignored-audience") # Audience is ignored for Azure DevOps + # Get token. + token = supplier.get_oidc_token("ignored-audience") # Audience is ignored for Azure DevOps. + + # Verify token result. if test_case.want_none: assert token is None else: assert token == test_case.want_token - # Verify the request was made correctly - if mock_post is not None: - expected_url = ( - "https://dev.azure.com/myorg/project-123/_apis/distributedtask/" - "hubs/build/plans/plan-456/jobs/job-789/oidctoken?api-version=7.2-preview.1" - ) - mock_post.assert_called_once_with( - expected_url, - headers={ - "Authorization": "Bearer azdo-access-token", - "Content-Type": "application/json", - "Content-Length": "0", - }, - ) + + # Verify the HTTP request was made correctly (only for successful token cases). + expected_url = ( + "https://dev.azure.com/myorg/project-123/_apis/distributedtask/" + "hubs/build/plans/plan-456/jobs/job-789/oidctoken?api-version=7.2-preview.1" + ) + mock_post.assert_called_once_with( + expected_url, + headers={ + "Authorization": "Bearer azure-devops-access-token", + "Content-Type": "application/json", + "Content-Length": "0", + }, + ) + + # For failure cases, verify HTTP request was still made but returned failure. + if test_case.want_none and test_case.response_ok is False: + mock_post.assert_called_once() # Request was made but failed. From 9f02727c4f1280cbb0c9c185aafc281b1f41a01f Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Mon, 29 Sep 2025 13:32:08 +0000 Subject: [PATCH 18/19] refactored code for github and azure devops oidc --- databricks/sdk/credentials_provider.py | 102 +++++++++++-------------- 1 file changed, 43 insertions(+), 59 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 5b43bd389..21399baea 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -12,7 +12,7 @@ import threading import time from datetime import datetime -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import google.auth # type: ignore import requests @@ -360,33 +360,47 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) -@oauth_credentials_strategy("github-oidc", ["host", "client_id"]) -def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]: +def _oidc_credentials_provider( + cfg: "Config", supplier_factory: Callable[[], Any], provider_name: str +) -> Optional[CredentialsProvider]: """ - DatabricksWIFCredentials uses a Token Supplier to get a JWT Token and exchanges - it for a Databricks Token. + Generic OIDC credentials provider that works with any OIDC token supplier. + + Args: + cfg: Databricks configuration + supplier_factory: Callable that returns an OIDC token supplier instance + provider_name: Human-readable name (e.g., "GitHub OIDC", "Azure DevOps OIDC") - Supported suppliers: - - GitHub OIDC + Returns: + OAuthCredentialsProvider if successful, None if supplier unavailable or token retrieval fails """ - supplier = oidc_token_supplier.GitHubOIDCTokenSupplier() + # Try to create the supplier + try: + supplier = supplier_factory() + except Exception as e: + logger.debug(f"{provider_name}: {str(e)}") + return None + # Determine the audience for token exchange audience = cfg.token_audience if audience is None and cfg.is_account_client: audience = cfg.account_id if audience is None and not cfg.is_account_client: audience = cfg.oidc_endpoints.token_endpoint - # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. + # Try to get an OIDC token. If no supplier returns a token, we cannot use this authentication mode. id_token = supplier.get_oidc_token(audience) if not id_token: + logger.debug(f"{provider_name}: no token available, skipping authentication method") return None + logger.info(f"Configured {provider_name} authentication") + def token_source_for(audience: str) -> oauth.TokenSource: id_token = supplier.get_oidc_token(audience) if not id_token: # Should not happen, since we checked it above. - raise Exception("Cannot get OIDC token") + raise Exception(f"Cannot get {provider_name} token") return oauth.ClientCredentials( client_id=cfg.client_id, @@ -412,6 +426,19 @@ def token() -> oauth.Token: return OAuthCredentialsProvider(refreshed_headers, token) +@oauth_credentials_strategy("github-oidc", ["host", "client_id"]) +def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]: + """ + GitHub OIDC authentication uses a Token Supplier to get a JWT Token and exchanges + it for a Databricks Token. + + Supported in GitHub Actions with OIDC service connections. + """ + return _oidc_credentials_provider( + cfg=cfg, supplier_factory=lambda: oidc_token_supplier.GitHubOIDCTokenSupplier(), provider_name="GitHub OIDC" + ) + + @oauth_credentials_strategy("azure-devops-oidc", ["host", "client_id"]) def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: """ @@ -420,54 +447,11 @@ def azure_devops_oidc(cfg: "Config") -> Optional[CredentialsProvider]: Supported in Azure DevOps pipelines with OIDC service connections. """ - try: - supplier = oidc_token_supplier.AzureDevOpsOIDCTokenSupplier() - except ValueError as e: - logger.debug(str(e)) - return None - - audience = cfg.token_audience - if audience is None and cfg.is_account_client: - audience = cfg.account_id - if audience is None and not cfg.is_account_client: - audience = cfg.oidc_endpoints.token_endpoint - - # Try to get an idToken. If no supplier returns a token, we cannot use this authentication mode. - id_token = supplier.get_oidc_token(audience) - if not id_token: - logger.debug("Azure DevOps OIDC: no token available, skipping authentication method") - return None - - logger.info("Configured Azure DevOps OIDC authentication") - - def token_source_for(audience: str) -> oauth.TokenSource: - id_token = supplier.get_oidc_token(audience) - if not id_token: - # Should not happen, since we checked it above. - raise Exception("Cannot get Azure DevOps OIDC token") - - return oauth.ClientCredentials( - client_id=cfg.client_id, - client_secret="", # we have no (rotatable) secrets in OIDC flow - token_url=cfg.oidc_endpoints.token_endpoint, - endpoint_params={ - "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", - "subject_token": id_token, - "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", - }, - scopes=["all-apis"], - use_params=True, - disable_async=cfg.disable_async_token_refresh, - ) - - def refreshed_headers() -> Dict[str, str]: - token = token_source_for(audience).token() - return {"Authorization": f"{token.token_type} {token.access_token}"} - - def token() -> oauth.Token: - return token_source_for(audience).token() - - return OAuthCredentialsProvider(refreshed_headers, token) + return _oidc_credentials_provider( + cfg=cfg, + supplier_factory=lambda: oidc_token_supplier.AzureDevOpsOIDCTokenSupplier(), + provider_name="Azure DevOps OIDC", + ) @oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"]) @@ -1078,10 +1062,10 @@ def __init__(self) -> None: env_oidc, file_oidc, github_oidc, - azure_devops_oidc, azure_service_principal, github_oidc_azure, azure_cli, + azure_devops_oidc, external_browser, databricks_cli, runtime_native_auth, From ac5db4f4c619ff76e8fd2d6f699fda15cb70e483 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Wed, 1 Oct 2025 16:40:51 +0000 Subject: [PATCH 19/19] addressed comments --- databricks/sdk/credentials_provider.py | 4 +++- databricks/sdk/oidc_token_supplier.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index 21399baea..022482370 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -435,7 +435,9 @@ def github_oidc(cfg: "Config") -> Optional[CredentialsProvider]: Supported in GitHub Actions with OIDC service connections. """ return _oidc_credentials_provider( - cfg=cfg, supplier_factory=lambda: oidc_token_supplier.GitHubOIDCTokenSupplier(), provider_name="GitHub OIDC" + cfg=cfg, + supplier_factory=lambda: oidc_token_supplier.GitHubOIDCTokenSupplier(), + provider_name="GitHub OIDC", ) diff --git a/databricks/sdk/oidc_token_supplier.py b/databricks/sdk/oidc_token_supplier.py index 70efd9d59..bd050dd5f 100644 --- a/databricks/sdk/oidc_token_supplier.py +++ b/databricks/sdk/oidc_token_supplier.py @@ -7,6 +7,7 @@ logger = logging.getLogger("databricks.sdk") +# TODO: Check the required environment variables while creating the instance rather than in the get_oidc_token method to allow early return. class GitHubOIDCTokenSupplier: """ Supplies OIDC tokens from GitHub Actions.