added Azure DevOps OIDC authentication

Author: Divyansh-db

databricks/sdk/credentials_provider.py

@@ -406,6 +406,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]:
@@ -1015,6 +1066,7 @@ def __init__(self) -> None:
             env_oidc,
             file_oidc,
             github_oidc,
+            azure_devops_oidc,
             azure_service_principal,
             github_oidc_azure,
             azure_cli,

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
+