added environment variables to config

Author: Divyansh-db

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")

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")

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]):