Support Databricks Workload Identity Federation for GitHub tokens (#423)

## What changes are proposed in this pull request?

This PR adds support for Databricks Workload Identity Federation using
GitHub tokens. This allows users to use WIF from their GitHub Workflows
and authenticate their workloads without long lived secrets.

This new credentials strategy is added to the DefaultCredentialsStrategy
after the other Databricks Credentials Strategy and before cloud
specific authentication methods.
WIF credentials uses a subset of configuration values of other
Databricks authentication methods. By being added after them it ensures
that WIF is not used when other Databricks authentication methods are
configured.
WIF uses the Databricks client id, which is not used by cloud specific
authentication methods. Therefore, it will not be used when cloud
specific authentication methods are configured.

## How is this tested?
Added tests.

Author: hectorcast-db

NEXT_CHANGELOG.md

@@ -3,6 +3,10 @@
 ## Release v0.48.0
 
 ### New Features and Improvements
+ * Introduce support for Databricks Workload Identity Federation in GitHub workflows ([423](https://github.com/databricks/databricks-sdk-java/pull/423)).
+   See README.md for instructions.
+ * [Breaking] Users running their workflows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST`
+   environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods.
 
 ### Bug Fixes
 

README.md

@@ -116,18 +116,18 @@ 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 basic (username/password) authentication (`auth_type="basic"` 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 Databricks Workload Identity Federation (WIF) authentication using OIDC (`auth_type="github-oidc"` argument).
 
 - For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents.
-- For Databricks basic authentication, you must provide `host`, `username`, and `password` _(for AWS workspace-level operations)_; or `host`, `account_id`, `username`, and `password` _(for AWS, Azure, or GCP account-level operations)_; 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.
 
 | Argument     | Description | Environment variable |
 |--------------|-------------|-------------------|
 | `host`       | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` |     
 | `account_id` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Only has effect when `Host` is either `https://accounts.cloud.databricks.com/` _(AWS)_, `https://accounts.azuredatabricks.net/` _(Azure)_, or `https://accounts.gcp.databricks.com/` _(GCP)_. | `DATABRICKS_ACCOUNT_ID` |
 | `token`      | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` |
-| `username`   | _(String)_ The Databricks username part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_USERNAME` |
-| `password`   | _(String)_ The Databricks password part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_PASSWORD` |
+| `client_id`      | _(String)_ The Databricks Service Principal Application ID.                                                                                                                                                                                                               | `DATABRICKS_CLIENT_ID`  |
+| `token_audience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier.                                                                                                                                         | `TOKEN_AUDIENCE`        |
 
 For example, to use Databricks token authentication:
 

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

@@ -141,6 +141,13 @@ public class DatabricksConfig {
 
   private DatabricksEnvironment databricksEnvironment;
 
+  /**
+   * When using Workload Identity Federation, the audience to specify when fetching an ID token from
+   * the ID token supplier.
+   */
+  @ConfigAttribute(env = "TOKEN_AUDIENCE")
+  private String tokenAudience;
+
   public Environment getEnv() {
     return env;
   }
@@ -512,6 +519,15 @@ public DatabricksConfig setHttpClient(HttpClient httpClient) {
     return this;
   }
 
+  public String getTokenAudience() {
+    return tokenAudience;
+  }
+
+  public DatabricksConfig setTokenAudience(String tokenAudience) {
+    this.tokenAudience = tokenAudience;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

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

@@ -1,9 +1,6 @@
 package com.databricks.sdk.core;
 
-import com.databricks.sdk.core.oauth.AzureGithubOidcCredentialsProvider;
-import com.databricks.sdk.core.oauth.AzureServicePrincipalCredentialsProvider;
-import com.databricks.sdk.core.oauth.ExternalBrowserCredentialsProvider;
-import com.databricks.sdk.core.oauth.OAuthM2MServicePrincipalCredentialsProvider;
+import com.databricks.sdk.core.oauth.*;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -18,6 +15,7 @@ public class DefaultCredentialsProvider implements CredentialsProvider {
           PatCredentialsProvider.class,
           BasicCredentialsProvider.class,
           OAuthM2MServicePrincipalCredentialsProvider.class,
+          GithubOidcCredentialsProvider.class,
           AzureGithubOidcCredentialsProvider.class,
           AzureServicePrincipalCredentialsProvider.class,
           AzureCliCredentialsProvider.class,

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

@@ -63,7 +63,7 @@ private static RefreshableTokenSource tokenSourceFor(DatabricksConfig config, St
         .withClientId(config.getAzureClientId())
         .withClientSecret(config.getAzureClientSecret())
         .withTokenUrl(tokenUrl)
-        .withEndpointParameters(endpointParams)
+        .withEndpointParametersSupplier(() -> endpointParams)
         .withAuthParameterPosition(AuthParameterPosition.BODY)
         .build();
   }

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

@@ -3,6 +3,7 @@
 import com.databricks.sdk.core.commons.CommonsHttpClient;
 import com.databricks.sdk.core.http.HttpClient;
 import java.util.*;
+import java.util.function.Supplier;
 
 /**
  * An implementation of RefreshableTokenSource implementing the client_credentials OAuth grant type.
@@ -18,7 +19,11 @@ public static class Builder {
     private String clientSecret;
     private String tokenUrl;
     private HttpClient hc = new CommonsHttpClient.Builder().withTimeoutSeconds(30).build();
-    private Map<String, String> endpointParams = Collections.emptyMap();
+
+    // Endpoint parameters can include tokens with expiration which
+    // may need to be refreshed. This supplier will be called each time
+    // the credentials are refreshed.
+    private Supplier<Map<String, String>> endpointParamsSupplier = null;
     private List<String> scopes = Collections.emptyList();
     private AuthParameterPosition position = AuthParameterPosition.BODY;
 
@@ -32,13 +37,14 @@ public Builder withClientSecret(String clientSecret) {
       return this;
     }
 
-    public Builder withTokenUrl(String tokenUrl) {
-      this.tokenUrl = tokenUrl;
+    public Builder withEndpointParametersSupplier(
+        Supplier<Map<String, String>> endpointParamsSupplier) {
+      this.endpointParamsSupplier = endpointParamsSupplier;
       return this;
     }
 
-    public Builder withEndpointParameters(Map<String, String> params) {
-      this.endpointParams = params;
+    public Builder withTokenUrl(String tokenUrl) {
+      this.tokenUrl = tokenUrl;
       return this;
     }
 
@@ -59,34 +65,33 @@ public Builder withHttpClient(HttpClient hc) {
 
     public ClientCredentials build() {
       Objects.requireNonNull(this.clientId, "clientId must be specified");
-      Objects.requireNonNull(this.clientSecret, "clientSecret must be specified");
       Objects.requireNonNull(this.tokenUrl, "tokenUrl must be specified");
       return new ClientCredentials(
-          hc, clientId, clientSecret, tokenUrl, endpointParams, scopes, position);
+          hc, clientId, clientSecret, tokenUrl, endpointParamsSupplier, scopes, position);
     }
   }
 
   private HttpClient hc;
   private String clientId;
   private String clientSecret;
   private String tokenUrl;
-  private Map<String, String> endpointParams;
   private List<String> scopes;
   private AuthParameterPosition position;
+  private Supplier<Map<String, String>> endpointParamsSupplier;
 
   private ClientCredentials(
       HttpClient hc,
       String clientId,
       String clientSecret,
       String tokenUrl,
-      Map<String, String> endpointParams,
+      Supplier<Map<String, String>> endpointParamsSupplier,
       List<String> scopes,
       AuthParameterPosition position) {
     this.hc = hc;
     this.clientId = clientId;
     this.clientSecret = clientSecret;
     this.tokenUrl = tokenUrl;
-    this.endpointParams = endpointParams;
+    this.endpointParamsSupplier = endpointParamsSupplier;
     this.scopes = scopes;
     this.position = position;
   }
@@ -98,8 +103,8 @@ protected Token refresh() {
     if (scopes != null) {
       params.put("scope", String.join(" ", scopes));
     }
-    if (endpointParams != null) {
-      params.putAll(endpointParams);
+    if (endpointParamsSupplier != null) {
+      params.putAll(endpointParamsSupplier.get());
     }
     return retrieveToken(hc, clientId, clientSecret, tokenUrl, params, new HashMap<>(), position);
   }

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

@@ -0,0 +1,79 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.http.HttpClient;
+import com.databricks.sdk.core.http.Request;
+import com.databricks.sdk.core.http.Response;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import java.io.IOException;
+
+public class GitHubOidcTokenSupplier {
+
+  private final ObjectMapper mapper = new ObjectMapper();
+  private final HttpClient httpClient;
+  private final String idTokenRequestUrl;
+  private final String idTokenRequestToken;
+  private final String tokenAudience;
+
+  public GitHubOidcTokenSupplier(
+      HttpClient httpClient,
+      String idTokenRequestUrl,
+      String idTokenRequestToken,
+      String tokenAudience) {
+    this.httpClient = httpClient;
+    this.idTokenRequestUrl = idTokenRequestUrl;
+    this.idTokenRequestToken = idTokenRequestToken;
+    this.tokenAudience = tokenAudience;
+  }
+
+  /** Checks if the required parameters are present to request a GitHub's OIDC token. */
+  public Boolean enabled() {
+    return idTokenRequestUrl != null && idTokenRequestToken != null;
+  }
+
+  /**
+   * Requests a GitHub's OIDC token.
+   *
+   * @return A GitHub OIDC token.
+   */
+  public String getOidcToken() {
+    if (!enabled()) {
+      throw new DatabricksException("Failed to request ID token: missing required parameters");
+    }
+
+    String requestUrl = idTokenRequestUrl;
+    if (tokenAudience != null) {
+      requestUrl += "&audience=" + tokenAudience;
+    }
+
+    Request req =
+        new Request("GET", requestUrl).withHeader("Authorization", "Bearer " + idTokenRequestToken);
+
+    Response resp;
+    try {
+      resp = httpClient.execute(req);
+    } catch (IOException e) {
+      throw new DatabricksException(
+          "Failed to request ID token from " + requestUrl + ":" + e.getMessage(), e);
+    }
+
+    if (resp.getStatusCode() != 200) {
+      throw new DatabricksException(
+          "Failed to request ID token: status code "
+              + resp.getStatusCode()
+              + ", response body: "
+              + resp.getBody().toString());
+    }
+
+    ObjectNode jsonResp;
+    try {
+      jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class);
+    } catch (IOException e) {
+      throw new DatabricksException(
+          "Failed to request ID token: corrupted token: " + e.getMessage());
+    }
+
+    return jsonResp.get("value").textValue();
+  }
+}

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

@@ -0,0 +1,67 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.CredentialsProvider;
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.HeaderFactory;
+import com.google.common.collect.ImmutableMap;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * GithubOidcCredentialsProvider uses a Token Supplier to get a GitHub OIDC JWT Token and exchanges
+ * it for a Databricks Token.
+ */
+public class GithubOidcCredentialsProvider implements CredentialsProvider {
+
+  @Override
+  public String authType() {
+    return "github-oidc";
+  }
+
+  @Override
+  public HeaderFactory configure(DatabricksConfig config) throws DatabricksException {
+    GitHubOidcTokenSupplier idTokenProvider =
+        new GitHubOidcTokenSupplier(
+            config.getHttpClient(),
+            config.getActionsIdTokenRequestUrl(),
+            config.getActionsIdTokenRequestToken(),
+            config.getTokenAudience());
+
+    if (!idTokenProvider.enabled() || config.getHost() == null || config.getClientId() == null) {
+      return null;
+    }
+
+    String endpointUrl;
+
+    try {
+      endpointUrl = config.getOidcEndpoints().getTokenEndpoint();
+    } catch (IOException e) {
+      throw new DatabricksException("Unable to fetch OIDC endpoint: " + e.getMessage(), e);
+    }
+
+    ClientCredentials clientCredentials =
+        new ClientCredentials.Builder()
+            .withHttpClient(config.getHttpClient())
+            .withClientId(config.getClientId())
+            .withTokenUrl(endpointUrl)
+            .withScopes(Collections.singletonList("all-apis"))
+            .withAuthParameterPosition(AuthParameterPosition.HEADER)
+            .withEndpointParametersSupplier(
+                () ->
+                    new ImmutableMap.Builder<String, String>()
+                        .put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt")
+                        .put("subject_token", idTokenProvider.getOidcToken())
+                        .put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange")
+                        .build())
+            .build();
+
+    return () -> {
+      Map<String, String> headers = new HashMap<>();
+      headers.put("Authorization", "Bearer " + clientCredentials.getToken().getAccessToken());
+      return headers;
+    };
+  }
+}