restructured code and tests

Author: Divyansh-db

databricks/sdk/credentials_provider.py

@@ -99,37 +99,21 @@ 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.
 
     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):
-                # 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,26 +412,19 @@ 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
     and exchanges it for a Databricks Token.
 
     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:

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