added tests, refactored code

Author: Divyansh-db

NEXT_CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### New Features and Improvements
 
+* Add native support for Azure DevOps OIDC authentication
+
 ### Bug Fixes
 
 ### Documentation

README.md

@@ -116,10 +116,11 @@ Depending on the Databricks authentication method, the SDK uses the following in
 
 ### Databricks native authentication
 
-By default, the Databricks SDK for Java 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 Java 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.
+- 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 |
 |--------------|-------------|-------------------|

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -431,13 +431,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) {
     return this;
   }
 
-  /** @deprecated Use {@link #getAzureUseMsi()} instead. */
+  /**
+   * @deprecated Use {@link #getAzureUseMsi()} instead.
+   */
   @Deprecated()
   public boolean getAzureUseMSI() {
     return azureUseMsi;
   }
 
-  /** @deprecated Use {@link #getAzureUseMsi()} instead. */
+  /**
+   * @deprecated Use {@link #getAzureUseMsi()} instead.
+   */
   @Deprecated
   public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) {
     this.azureUseMsi = azureUseMsi;

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java

@@ -131,8 +131,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
     try {
       namedIdTokenSources.add(
           new NamedIDTokenSource(
-              "azure-devops-oidc",
-              new AzureDevOpsIDTokenSource(config.getHttpClient())));
+              "azure-devops-oidc", new AzureDevOpsIDTokenSource(config.getHttpClient())));
     } catch (DatabricksException e) {
       LOG.debug("Azure DevOps OIDC provider not available: {}", e.getMessage());
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSource.java

@@ -4,6 +4,7 @@
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
 import com.databricks.sdk.core.http.Response;
+import com.databricks.sdk.core.utils.Environment;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.base.Strings;
@@ -13,7 +14,7 @@
  * AzureDevOpsIDTokenSource retrieves JWT Tokens from Azure DevOps Pipelines. This class implements
  * the IDTokenSource interface and provides a method for obtaining ID tokens specifically from Azure
  * DevOps Pipeline environment.
- * 
+ *
  * <p>This implementation follows the Azure DevOps OIDC token API as documented at:
  * https://learn.microsoft.com/en-us/rest/api/azure/devops/distributedtask/oidctoken/create
  */
@@ -32,32 +33,56 @@ public class AzureDevOpsIDTokenSource implements IDTokenSource {
   private final String azureDevOpsHostType;
   /* HTTP client for making requests to Azure DevOps */
   private final HttpClient httpClient;
+  /* Environment for reading configuration values */
+  private final Environment environment;
   /* JSON mapper for parsing response data */
   private static final ObjectMapper mapper = new ObjectMapper();
 
   /**
-   * Constructs a new AzureDevOpsIDTokenSource by reading environment variables.
-   * This constructor implements fail-early validation - if any required environment
-   * variables are missing, it will throw a DatabricksException immediately.
+   * Constructs a new AzureDevOpsIDTokenSource by reading environment variables. This constructor
+   * implements fail-early validation - if any required environment variables are missing, it will
+   * throw a DatabricksException immediately.
    *
    * @param httpClient The HTTP client to use for making requests
    * @throws DatabricksException if any required environment variables are missing
    */
   public AzureDevOpsIDTokenSource(HttpClient httpClient) {
+    this(httpClient, createDefaultEnvironment());
+  }
+
+  /**
+   * Constructs a new AzureDevOpsIDTokenSource with a custom environment. This constructor is
+   * primarily used for testing to inject mock environment variables.
+   *
+   * @param httpClient The HTTP client to use for making requests
+   * @param environment The environment to read configuration from
+   * @throws DatabricksException if httpClient is null or any required environment variables are
+   *     missing
+   */
+  public AzureDevOpsIDTokenSource(HttpClient httpClient, Environment environment) {
     if (httpClient == null) {
       throw new DatabricksException("HttpClient cannot be null");
     }
     this.httpClient = httpClient;
+    this.environment = environment;
 
-    // Fail early: validate all required environment variables
     this.azureDevOpsAccessToken = validateEnvironmentVariable("SYSTEM_ACCESSTOKEN");
-    this.azureDevOpsTeamFoundationCollectionUri = validateEnvironmentVariable("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI");
+    this.azureDevOpsTeamFoundationCollectionUri =
+        validateEnvironmentVariable("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI");
     this.azureDevOpsPlanId = validateEnvironmentVariable("SYSTEM_PLANID");
     this.azureDevOpsJobId = validateEnvironmentVariable("SYSTEM_JOBID");
     this.azureDevOpsTeamProjectId = validateEnvironmentVariable("SYSTEM_TEAMPROJECTID");
     this.azureDevOpsHostType = validateEnvironmentVariable("SYSTEM_HOSTTYPE");
   }
 
+  /** Creates a default Environment using system environment variables. */
+  private static Environment createDefaultEnvironment() {
+    String pathEnv = System.getenv("PATH");
+    String[] pathArray =
+        pathEnv != null ? pathEnv.split(java.io.File.pathSeparator) : new String[0];
+    return new Environment(System.getenv(), pathArray, System.getProperty("os.name"));
+  }
+
   /**
    * Validates that an environment variable is present and not empty.
    *
@@ -66,8 +91,14 @@ public AzureDevOpsIDTokenSource(HttpClient httpClient) {
    * @throws DatabricksException if the environment variable is missing or empty
    */
   private String validateEnvironmentVariable(String varName) {
-    String value = System.getenv(varName);
+    String value = environment.get(varName);
     if (Strings.isNullOrEmpty(value)) {
+      if (varName.equals("SYSTEM_ACCESSTOKEN")) {
+        throw new DatabricksException(
+            String.format(
+                "Missing %s, 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",
+                varName));
+      }
       throw new DatabricksException(
           String.format("Missing %s, likely not calling from Azure DevOps Pipeline", varName));
     }
@@ -76,9 +107,9 @@ private String validateEnvironmentVariable(String varName) {
 
   /**
    * Retrieves an ID token from Azure DevOps Pipelines. This method makes an authenticated request
-   * to Azure DevOps to obtain a JWT token that can later be exchanged for a Databricks access token.
+   * to Azure DevOps to obtain a JWT token that can later be exchanged for a Databricks access
+   * token.
    *
-   * 
    * <p>Note: The audience parameter is ignored for Azure DevOps OIDC tokens as they have a
    * hardcoded audience for Azure AD integration.
    *
@@ -89,26 +120,30 @@ private String validateEnvironmentVariable(String varName) {
   @Override
   public IDToken getIDToken(String audience) {
 
-    // Build Azure DevOps OIDC endpoint URL
-    // Format: {collectionUri}/{teamProjectId}/_apis/distributedtask/hubs/{hostType}/plans/{planId}/jobs/{jobId}/oidctoken?api-version=7.2-preview.1
-    String requestUrl = String.format(
-        "%s/%s/_apis/distributedtask/hubs/%s/plans/%s/jobs/%s/oidctoken?api-version=7.2-preview.1",
-        azureDevOpsTeamFoundationCollectionUri,
-        azureDevOpsTeamProjectId,
-        azureDevOpsHostType,
-        azureDevOpsPlanId,
-        azureDevOpsJobId);
+    // Build Azure DevOps OIDC endpoint URL.
+    // Format:
+    // {collectionUri}/{teamProjectId}/_apis/distributedtask/hubs/{hostType}/plans/{planId}/jobs/{jobId}/oidctoken?api-version=7.2-preview.1
+    String requestUrl =
+        String.format(
+            "%s/%s/_apis/distributedtask/hubs/%s/plans/%s/jobs/%s/oidctoken?api-version=7.2-preview.1",
+            azureDevOpsTeamFoundationCollectionUri,
+            azureDevOpsTeamProjectId,
+            azureDevOpsHostType,
+            azureDevOpsPlanId,
+            azureDevOpsJobId);
 
-    Request req = new Request("POST", requestUrl)
-        .withHeader("Authorization", "Bearer " + azureDevOpsAccessToken)
-        .withHeader("Content-Type", "application/json");
+    Request req =
+        new Request("POST", requestUrl)
+            .withHeader("Authorization", "Bearer " + azureDevOpsAccessToken)
+            .withHeader("Content-Type", "application/json");
 
     Response resp;
     try {
       resp = httpClient.execute(req);
     } catch (IOException e) {
       throw new DatabricksException(
-          "Failed to request ID token from Azure DevOps at " + requestUrl + ": " + e.getMessage(), e);
+          "Failed to request ID token from Azure DevOps at " + requestUrl + ": " + e.getMessage(),
+          e);
     }
 
     if (resp.getStatusCode() != 200) {
@@ -129,7 +164,6 @@ public IDToken getIDToken(String audience) {
           "Failed to parse Azure DevOps OIDC token response: " + e.getMessage(), e);
     }
 
-    // Validate response structure and token value
     if (!jsonResp.has("oidcToken")) {
       throw new DatabricksException("Azure DevOps OIDC token response missing 'oidcToken' field");
     }
@@ -141,7 +175,8 @@ public IDToken getIDToken(String audience) {
       }
       return new IDToken(tokenValue);
     } catch (IllegalArgumentException e) {
-      throw new DatabricksException("Received invalid OIDC token from Azure DevOps: " + e.getMessage(), e);
+      throw new DatabricksException(
+          "Received invalid OIDC token from Azure DevOps: " + e.getMessage(), e);
     }
   }
 }