Emmyzhou db/direct dataplane access

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

@@ -34,14 +34,6 @@ public NamedIDTokenSource(String name, IDTokenSource idTokenSource) {
       this.name = name;
       this.idTokenSource = idTokenSource;
     }
-
-    public String getName() {
-      return name;
-    }
-
-    public IDTokenSource getIdTokenSource() {
-      return idTokenSource;
-    }
   }
 
   public DefaultCredentialsProvider() {}
@@ -143,14 +135,13 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
                   config.getClientId(),
                   config.getHost(),
                   endpoints,
-                  namedIdTokenSource.getIdTokenSource(),
+                  namedIdTokenSource.idTokenSource,
                   config.getHttpClient())
               .audience(config.getTokenAudience())
               .accountId(config.isAccountClient() ? config.getAccountId() : null)
               .build();
 
-      providers.add(
-          new TokenSourceCredentialsProvider(oauthTokenSource, namedIdTokenSource.getName()));
+      providers.add(new TokenSourceCredentialsProvider(oauthTokenSource, namedIdTokenSource.name));
     }
   }
 

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

@@ -0,0 +1,102 @@
+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 DatabricksOAuthTokenSource cpTokenSource;
+  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 DatabricksOAuthTokenSource} for control plane tokens.
+   * @throws NullPointerException if either parameter is null
+   */
+  public DataPlaneTokenSource(HttpClient httpClient, DatabricksOAuthTokenSource cpTokenSource) {
+    this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
+    this.cpTokenSource =
+        Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
+    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
+   */
+  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));
+
+    return specificSource.getToken();
+  }
+}

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

@@ -1,15 +1,12 @@
 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 com.google.common.base.Strings;
-import java.io.IOException;
 import java.time.LocalDateTime;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -44,8 +41,6 @@ public class DatabricksOAuthTokenSource extends RefreshableTokenSource {
   private static final String SCOPE_PARAM = "scope";
   private static final String CLIENT_ID_PARAM = "client_id";
 
-  private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-
   private DatabricksOAuthTokenSource(Builder builder) {
     this.clientId = builder.clientId;
     this.host = builder.host;
@@ -123,44 +118,29 @@ public DatabricksOAuthTokenSource build() {
     }
   }
 
-  /**
-   * Validates that a value is non-null for required fields. If the value is a string, it also
-   * checks that it is non-empty.
-   *
-   * @param value The value to validate.
-   * @param fieldName The name of the field being validated.
-   * @throws IllegalArgumentException when the value is null or an empty string.
-   */
-  private static void validate(Object value, String fieldName) {
-    if (value == null) {
-      LOG.error("Required parameter '{}' is null", fieldName);
-      throw new IllegalArgumentException(
-          String.format("Required parameter '%s' cannot be null", fieldName));
-    }
-    if (value instanceof String && ((String) value).isEmpty()) {
-      LOG.error("Required parameter '{}' is empty", fieldName);
-      throw new IllegalArgumentException(
-          String.format("Required parameter '%s' cannot be empty", fieldName));
-    }
-  }
-
   /**
    * Retrieves an OAuth token by exchanging an ID token. Implements the OAuth token exchange flow to
    * obtain an access token.
    *
    * @return A Token containing the access token and related information.
    * @throws DatabricksException when the token exchange fails.
-   * @throws IllegalArgumentException when there is an error code in the response or when required
-   *     parameters are missing.
+   * @throws IllegalArgumentException when the required string parameters are empty.
+   * @throws NullPointerException when any of the required parameters are null.
    */
   @Override
   public Token refresh() {
-    // Validate all required parameters
-    validate(clientId, "ClientID");
-    validate(host, "Host");
-    validate(endpoints, "Endpoints");
-    validate(idTokenSource, "IDTokenSource");
-    validate(httpClient, "HttpClient");
+    Objects.requireNonNull(clientId, "ClientID cannot be null");
+    Objects.requireNonNull(host, "Host cannot be null");
+    Objects.requireNonNull(endpoints, "Endpoints cannot be null");
+    Objects.requireNonNull(idTokenSource, "IDTokenSource cannot be null");
+    Objects.requireNonNull(httpClient, "HttpClient cannot be null");
+
+    if (clientId.isEmpty()) {
+      throw new IllegalArgumentException("ClientID cannot be empty");
+    }
+    if (host.isEmpty()) {
+      throw new IllegalArgumentException("Host cannot be empty");
+    }
 
     String effectiveAudience = determineAudience();
     IDToken idToken = idTokenSource.getIDToken(effectiveAudience);
@@ -172,47 +152,20 @@ public Token refresh() {
     params.put(SCOPE_PARAM, SCOPE);
     params.put(CLIENT_ID_PARAM, clientId);
 
-    Response rawResponse;
-    try {
-      rawResponse = httpClient.execute(new FormRequest(endpoints.getTokenEndpoint(), params));
-    } catch (IOException e) {
-      LOG.error(
-          "Failed to exchange ID token for access token at {}: {}",
-          endpoints.getTokenEndpoint(),
-          e.getMessage(),
-          e);
-      throw new DatabricksException(
-          String.format(
-              "Failed to exchange ID token for access token at %s: %s",
-              endpoints.getTokenEndpoint(), e.getMessage()),
-          e);
-    }
-
     OAuthResponse response;
     try {
-      response = OBJECT_MAPPER.readValue(rawResponse.getBody(), OAuthResponse.class);
-    } catch (IOException e) {
+      response =
+          TokenEndpointClient.requestToken(this.httpClient, endpoints.getTokenEndpoint(), params);
+    } catch (DatabricksException e) {
       LOG.error(
-          "Failed to parse OAuth response from token endpoint {}: {}",
+          "OAuth token exchange failed for client ID '{}' at {}: {}",
+          this.clientId,
           endpoints.getTokenEndpoint(),
           e.getMessage(),
           e);
-      throw new DatabricksException(
-          String.format(
-              "Failed to parse OAuth response from token endpoint %s: %s",
-              endpoints.getTokenEndpoint(), e.getMessage()));
+      throw e;
     }
 
-    if (response.getErrorCode() != null) {
-      LOG.error(
-          "Token exchange failed with error: {} - {}",
-          response.getErrorCode(),
-          response.getErrorSummary());
-      throw new IllegalArgumentException(
-          String.format(
-              "Token exchange failed with error: %s - %s",
-              response.getErrorCode(), response.getErrorSummary()));
-    }
     LocalDateTime expiry = LocalDateTime.now().plusSeconds(response.getExpiresIn());
     return new Token(
         response.getAccessToken(), response.getTokenType(), response.getRefreshToken(), expiry);

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

@@ -0,0 +1,92 @@
+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 DatabricksOAuthTokenSource} 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 DatabricksOAuthTokenSource cpTokenSource;
+  private final String authDetails;
+  private final HttpClient httpClient;
+
+  /**
+   * Constructs a new EndpointTokenSource.
+   *
+   * @param cpTokenSource The {@link DatabricksOAuthTokenSource} 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.
+   * @throws IllegalArgumentException if authDetails is empty.
+   * @throws NullPointerException if any of the parameters are null.
+   */
+  public EndpointTokenSource(
+      DatabricksOAuthTokenSource cpTokenSource, String authDetails, HttpClient httpClient) {
+    this.cpTokenSource =
+        Objects.requireNonNull(cpTokenSource, "Control plane token source cannot be null");
+    this.authDetails = Objects.requireNonNull(authDetails, "Authorization details cannot be null");
+    if (authDetails.isEmpty()) {
+      throw new IllegalArgumentException("Authorization details cannot be empty");
+    }
+    this.httpClient = Objects.requireNonNull(httpClient, "HTTP client cannot be null");
+  }
+
+  /**
+   * 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, 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);
+  }
+}