Add DataPlaneTokenSource and EndpointTokenSource

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

@@ -0,0 +1,115 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.http.HttpClient;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Manages and provides Databricks data plane tokens. This class is responsible for acquiring and
+ * caching OAuth tokens that are specific to a particular Databricks data plane service endpoint and
+ * a set of authorization details. It utilizes a {@link DatabricksOAuthTokenSource} for obtaining
+ * control plane tokens, which may then be exchanged or used to authorize requests for data plane
+ * tokens. Cached {@link EndpointTokenSource} instances are used to efficiently reuse tokens for
+ * repeated requests to the same endpoint with the same authorization context.
+ */
+public class DataPlaneTokenSource {
+  private final HttpClient httpClient;
+  private final TokenSource cpTokenSource;
+  private final String host;
+  private final ConcurrentHashMap<TokenSourceKey, EndpointTokenSource> sourcesCache;
+  /**
+   * Caching key for {@link EndpointTokenSource}, based on endpoint and authorization details. This
+   * is a value object that uniquely identifies a token source configuration.
+   */
+  private static final class TokenSourceKey {
+    /** The target service endpoint URL. */
+    private final String endpoint;
+
+    /** Specific authorization details for the endpoint. */
+    private final String authDetails;
+
+    /**
+     * Constructs a TokenSourceKey.
+     *
+     * @param endpoint The target service endpoint URL.
+     * @param authDetails Specific authorization details.
+     */
+    public TokenSourceKey(String endpoint, String authDetails) {
+      this.endpoint = endpoint;
+      this.authDetails = authDetails;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      TokenSourceKey that = (TokenSourceKey) o;
+      return Objects.equals(endpoint, that.endpoint)
+          && Objects.equals(authDetails, that.authDetails);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(endpoint, authDetails);
+    }
+  }
+
+  /**
+   * Constructs a DataPlaneTokenSource.
+   *
+   * @param httpClient The {@link HttpClient} for token requests.
+   * @param cpTokenSource The {@link TokenSource} for control plane tokens.
+   * @param host The host for the token exchange request.
+   * @throws NullPointerException if any parameter is null.
+   * @throws IllegalArgumentException if the host is empty.
+   */
+  public DataPlaneTokenSource(HttpClient httpClient, TokenSource cpTokenSource, String host) {
+    this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
+    this.cpTokenSource =
+        Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
+    this.host = Objects.requireNonNull(host, "Host cannot be null");
+
+    if (host.isEmpty()) {
+      throw new IllegalArgumentException("Host cannot be empty");
+    }
+    this.sourcesCache = new ConcurrentHashMap<>();
+  }
+
+  /**
+   * Retrieves a token for the specified endpoint and authorization details. It uses a cached {@link
+   * EndpointTokenSource} if available, otherwise creates and caches a new one.
+   *
+   * @param endpoint The target data plane service endpoint.
+   * @param authDetails Authorization details for the endpoint.
+   * @return The dataplane {@link Token}.
+   * @throws NullPointerException if either parameter is null.
+   * @throws IllegalArgumentException if either parameter is empty.
+   * @throws DatabricksException if the token request fails.
+   */
+  public Token getToken(String endpoint, String authDetails) {
+    Objects.requireNonNull(endpoint, "Data plane endpoint URL cannot be null");
+    Objects.requireNonNull(authDetails, "Authorization details cannot be null");
+
+    if (endpoint.isEmpty()) {
+      throw new IllegalArgumentException("Data plane endpoint URL cannot be empty");
+    }
+    if (authDetails.isEmpty()) {
+      throw new IllegalArgumentException("Authorization details cannot be empty");
+    }
+
+    TokenSourceKey key = new TokenSourceKey(endpoint, authDetails);
+
+    EndpointTokenSource specificSource =
+        sourcesCache.computeIfAbsent(
+            key,
+            k ->
+                new EndpointTokenSource(
+                    this.cpTokenSource, k.authDetails, this.httpClient, this.host));
+
+    return specificSource.getToken();
+  }
+}

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

@@ -0,0 +1,97 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.http.HttpClient;
+import java.time.LocalDateTime;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Represents a token source that exchanges a control plane token for an endpoint-specific dataplane
+ * token. It utilizes an underlying {@link TokenSource} to obtain the initial control plane token.
+ */
+public class EndpointTokenSource extends RefreshableTokenSource {
+  private static final Logger LOG = LoggerFactory.getLogger(EndpointTokenSource.class);
+  private static final String JWT_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer";
+  private static final String GRANT_TYPE_PARAM = "grant_type";
+  private static final String AUTHORIZATION_DETAILS_PARAM = "authorization_details";
+  private static final String ASSERTION_PARAM = "assertion";
+  private static final String TOKEN_ENDPOINT = "/oidc/v1/token";
+
+  private final TokenSource cpTokenSource;
+  private final String authDetails;
+  private final HttpClient httpClient;
+  private final String host;
+
+  /**
+   * Constructs a new EndpointTokenSource.
+   *
+   * @param cpTokenSource The {@link TokenSource} used to obtain the control plane token.
+   * @param authDetails The authorization details required for the token exchange.
+   * @param httpClient The {@link HttpClient} used to make the token exchange request.
+   * @param host The host for the token exchange request.
+   * @throws IllegalArgumentException if authDetails is empty or host is empty.
+   * @throws NullPointerException if any of the parameters are null.
+   */
+  public EndpointTokenSource(
+      TokenSource cpTokenSource, String authDetails, HttpClient httpClient, String host) {
+    this.cpTokenSource =
+        Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
+    this.authDetails = Objects.requireNonNull(authDetails, "Authorization details cannot be null");
+    this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
+    this.host = Objects.requireNonNull(host, "Host cannot be null");
+
+    if (authDetails.isEmpty()) {
+      throw new IllegalArgumentException("Authorization details cannot be empty");
+    }
+    if (host.isEmpty()) {
+      throw new IllegalArgumentException("Host cannot be empty");
+    }
+  }
+
+  /**
+   * Fetches an endpoint-specific dataplane token by exchanging a control plane token.
+   *
+   * <p>This method first obtains a control plane token from the configured {@code cpTokenSource}.
+   * It then uses this token as an assertion along with the provided {@code authDetails} to request
+   * a new, more scoped dataplane token from the Databricks OAuth token endpoint ({@value
+   * #TOKEN_ENDPOINT}).
+   *
+   * @return A new {@link Token} containing the exchanged dataplane access token, its type, any
+   *     accompanying refresh token, and its expiry time.
+   * @throws DatabricksException if the token exchange with the OAuth endpoint fails.
+   * @throws IllegalArgumentException if the token endpoint url is empty.
+   * @throws NullPointerException if any of the parameters are null.
+   */
+  @Override
+  protected Token refresh() {
+    Token cpToken = cpTokenSource.getToken();
+    Map<String, String> params = new HashMap<>();
+    params.put(GRANT_TYPE_PARAM, JWT_GRANT_TYPE);
+    params.put(AUTHORIZATION_DETAILS_PARAM, authDetails);
+    params.put(ASSERTION_PARAM, cpToken.getAccessToken());
+
+    OAuthResponse oauthResponse;
+    try {
+      oauthResponse =
+          TokenEndpointClient.requestToken(this.httpClient, this.host + TOKEN_ENDPOINT, params);
+    } catch (DatabricksException | IllegalArgumentException | NullPointerException e) {
+      LOG.error(
+          "Failed to exchange control plane token for dataplane token at endpoint {}: {}",
+          TOKEN_ENDPOINT,
+          e.getMessage(),
+          e);
+      throw e;
+    }
+
+    LocalDateTime expiry = LocalDateTime.now().plusSeconds(oauthResponse.getExpiresIn());
+    return new Token(
+        oauthResponse.getAccessToken(),
+        oauthResponse.getTokenType(),
+        oauthResponse.getRefreshToken(),
+        expiry);
+  }
+}

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

@@ -0,0 +1,91 @@
+package com.databricks.sdk.core.oauth;
+
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.http.FormRequest;
+import com.databricks.sdk.core.http.HttpClient;
+import com.databricks.sdk.core.http.Response;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.util.Map;
+import java.util.Objects;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Client for interacting with an OAuth token endpoint.
+ *
+ * <p>This class provides a method to request an OAuth token from a specified token endpoint URL
+ * using the provided HTTP client and request parameters. It handles the HTTP request and parses the
+ * JSON response into an {@link OAuthResponse} object.
+ */
+public final class TokenEndpointClient {
+  private static final Logger LOG = LoggerFactory.getLogger(TokenEndpointClient.class);
+  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+  private TokenEndpointClient() {}
+
+  /**
+   * Requests an OAuth token from the specified token endpoint.
+   *
+   * @param httpClient The {@link HttpClient} to use for making the request.
+   * @param tokenEndpointUrl The URL of the token endpoint.
+   * @param params A map of parameters to include in the token request.
+   * @return An {@link OAuthResponse} containing the token information.
+   * @throws DatabricksException if an error occurs during the token request or response parsing.
+   * @throws IllegalArgumentException if the token endpoint URL is empty.
+   * @throws NullPointerException if any of the parameters are null.
+   */
+  public static OAuthResponse requestToken(
+      HttpClient httpClient, String tokenEndpointUrl, Map<String, String> params)
+      throws DatabricksException {
+    Objects.requireNonNull(httpClient, "HttpClient cannot be null");
+    Objects.requireNonNull(params, "Request parameters map cannot be null");
+    Objects.requireNonNull(tokenEndpointUrl, "Token endpoint URL cannot be null");
+
+    if (tokenEndpointUrl.isEmpty()) {
+      throw new IllegalArgumentException("Token endpoint URL cannot be empty");
+    }
+
+    Response rawResponse;
+    try {
+      LOG.debug("Requesting token from endpoint: {}", tokenEndpointUrl);
+      rawResponse = httpClient.execute(new FormRequest(tokenEndpointUrl, params));
+    } catch (IOException e) {
+      LOG.error("Failed to request token from {}: {}", tokenEndpointUrl, e.getMessage(), e);
+      throw new DatabricksException(
+          String.format("Failed to request token from %s: %s", tokenEndpointUrl, e.getMessage()),
+          e);
+    }
+
+    OAuthResponse response;
+    try {
+      response = OBJECT_MAPPER.readValue(rawResponse.getBody(), OAuthResponse.class);
+    } catch (IOException e) {
+      LOG.error(
+          "Failed to parse OAuth response from token endpoint {}: {}",
+          tokenEndpointUrl,
+          e.getMessage(),
+          e);
+      throw new DatabricksException(
+          String.format(
+              "Failed to parse OAuth response from token endpoint %s: %s",
+              tokenEndpointUrl, e.getMessage()),
+          e);
+    }
+
+    if (response.getErrorCode() != null) {
+      String errorSummary =
+          response.getErrorSummary() != null ? response.getErrorSummary() : "No summary provided.";
+      LOG.error(
+          "Token request to {} failed with error: {} - {}",
+          tokenEndpointUrl,
+          response.getErrorCode(),
+          errorSummary);
+      throw new DatabricksException(
+          String.format(
+              "Token request failed with error: %s - %s", response.getErrorCode(), errorSummary));
+    }
+    LOG.debug("Successfully obtained token response from {}", tokenEndpointUrl);
+    return response;
+  }
+}