Add support for direct dataplane access

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.52.0
 
 ### New Features and Improvements
+* Added Direct-to-Dataplane API support, allowing users to query route optimized model serving endpoints ([#453](https://github.com/databricks/databricks-sdk-java/pull/453)).
 
 ### Bug Fixes
 

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

@@ -4,6 +4,7 @@
 import com.databricks.sdk.core.error.PrivateLinkInfo;
 import com.databricks.sdk.core.http.HttpClient;
 import com.databricks.sdk.core.http.Request;
+import com.databricks.sdk.core.http.RequestOptions;
 import com.databricks.sdk.core.http.Response;
 import com.databricks.sdk.core.retry.RequestBasedRetryStrategyPicker;
 import com.databricks.sdk.core.retry.RetryStrategy;
@@ -50,7 +51,6 @@ public Builder withDatabricksConfig(DatabricksConfig config) {
       this.accountId = config.getAccountId();
       this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost());
       this.isDebugHeaders = config.isDebugHeaders();
-
       return this;
     }
 
@@ -173,33 +173,42 @@ public Map<String, String> getStringMap(Request req) {
 
   protected <I, O> O withJavaType(Request request, JavaType javaType) {
     try {
-      Response response = getResponse(request);
+      Response response = executeInner(request, request.getUrl(), new RequestOptions());
       return deserialize(response.getBody(), javaType);
     } catch (IOException e) {
       throw new DatabricksException("IO error: " + e.getMessage(), e);
     }
   }
 
   /**
-   * Executes HTTP request with retries and converts it to proper POJO
+   * Executes HTTP request with retries and converts it to proper POJO.
    *
    * @param in Commons HTTP request
    * @param target Expected pojo type
    * @return POJO of requested type
    */
   public <T> T execute(Request in, Class<T> target) throws IOException {
-    Response out = getResponse(in);
+    return execute(in, target, new RequestOptions());
+  }
+
+  /**
+   * Executes HTTP request with retries and converts it to proper POJO, using custom request
+   * options.
+   *
+   * @param in Commons HTTP request
+   * @param target Expected pojo type
+   * @param options Optional request options to customize request behavior
+   * @return POJO of requested type
+   */
+  public <T> T execute(Request in, Class<T> target, RequestOptions options) throws IOException {
+    Response out = executeInner(in, in.getUrl(), options);
     if (target == Void.class) {
       return null;
     }
     return deserialize(out, target);
   }
 
-  private Response getResponse(Request in) {
-    return executeInner(in, in.getUrl());
-  }
-
-  private Response executeInner(Request in, String path) {
+  private Response executeInner(Request in, String path, RequestOptions options) {
     RetryStrategy retryStrategy = retryStrategyPicker.getRetryStrategy(in);
     int attemptNumber = 0;
     while (true) {
@@ -224,6 +233,8 @@ private Response executeInner(Request in, String path) {
       }
       in.withHeader("User-Agent", userAgent);
 
+      options.applyOptions(in);
+
       // Make the request, catching any exceptions, as we may want to retry.
       try {
         out = httpClient.execute(in);
@@ -434,4 +445,8 @@ public String serialize(Object body) throws JsonProcessingException {
     }
     return mapper.writeValueAsString(body);
   }
+
+  public HttpClient getHttpClient() {
+    return httpClient;
+  }
 }

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

@@ -1,5 +1,6 @@
 package com.databricks.sdk.core;
 
+import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.oauth.Token;
 import com.databricks.sdk.core.utils.AzureUtils;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -68,7 +69,7 @@ private Optional<String> getSubscription(DatabricksConfig config) {
   }
 
   @Override
-  public HeaderFactory configure(DatabricksConfig config) {
+  public OAuthHeaderFactory configure(DatabricksConfig config) {
     if (!config.isAzure()) {
       return null;
     }
@@ -86,15 +87,17 @@ public HeaderFactory configure(DatabricksConfig config) {
         mgmtTokenSource = null;
       }
       CliTokenSource finalMgmtTokenSource = mgmtTokenSource;
-      return () -> {
-        Token token = tokenSource.getToken();
-        Map<String, String> headers = new HashMap<>();
-        headers.put("Authorization", token.getTokenType() + " " + token.getAccessToken());
-        if (finalMgmtTokenSource != null) {
-          AzureUtils.addSpManagementToken(finalMgmtTokenSource, headers);
-        }
-        return AzureUtils.addWorkspaceResourceId(config, headers);
-      };
+      return OAuthHeaderFactory.fromSuppliers(
+          tokenSource::getToken,
+          () -> {
+            Token token = tokenSource.getToken();
+            Map<String, String> headers = new HashMap<>();
+            headers.put("Authorization", token.getTokenType() + " " + token.getAccessToken());
+            if (finalMgmtTokenSource != null) {
+              AzureUtils.addSpManagementToken(finalMgmtTokenSource, headers);
+            }
+            return AzureUtils.addWorkspaceResourceId(config, headers);
+          });
     } catch (DatabricksException e) {
       String stderr = e.getMessage();
       if (stderr.contains("not found")) {

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

@@ -1,6 +1,6 @@
 package com.databricks.sdk.core;
 
-import com.databricks.sdk.core.oauth.Token;
+import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.utils.OSUtils;
 import java.util.*;
 import org.slf4j.Logger;
@@ -36,7 +36,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
   }
 
   @Override
-  public HeaderFactory configure(DatabricksConfig config) {
+  public OAuthHeaderFactory configure(DatabricksConfig config) {
     String host = config.getHost();
     if (host == null) {
       return null;
@@ -48,12 +48,7 @@ public HeaderFactory configure(DatabricksConfig config) {
         return null;
       }
       tokenSource.getToken(); // We need this for checking if databricks CLI is installed.
-      return () -> {
-        Token token = tokenSource.getToken();
-        Map<String, String> headers = new HashMap<>();
-        headers.put("Authorization", token.getTokenType() + " " + token.getAccessToken());
-        return headers;
-      };
+      return OAuthHeaderFactory.fromTokenSource(tokenSource);
     } catch (DatabricksException e) {
       String stderr = e.getMessage();
       if (stderr.contains("not found")) {

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

@@ -4,7 +4,10 @@
 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.oauth.ErrorTokenSource;
+import com.databricks.sdk.core.oauth.OAuthHeaderFactory;
 import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.Cloud;
 import com.databricks.sdk.core.utils.Environment;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -209,6 +212,24 @@ public synchronized Map<String, String> authenticate() throws DatabricksExceptio
     }
   }
 
+  public TokenSource getTokenSource() {
+    if (headerFactory == null) {
+      try {
+        ConfigLoader.fixHostIfNeeded(this);
+        headerFactory = credentialsProvider.configure(this);
+      } catch (Exception e) {
+        return new ErrorTokenSource("Failed to get token source: " + e.getMessage());
+      }
+      setAuthType(credentialsProvider.authType());
+    }
+
+    if (headerFactory instanceof OAuthHeaderFactory) {
+      return (TokenSource) headerFactory;
+    }
+    return new ErrorTokenSource(
+        String.format("OAuth Token not supported for current auth type %s", authType));
+  }
+
   public CredentialsProvider getCredentialsProvider() {
     return this.credentialsProvider;
   }

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/http/RequestOptions.java

@@ -0,0 +1,71 @@
+package com.databricks.sdk.core.http;
+
+import java.util.function.Function;
+
+/**
+ * A builder class for configuring HTTP request transformations including authentication, URL, and
+ * user agent headers.
+ *
+ * <p>Experimental: this class is experimental and subject to change in backward incompatible ways.
+ */
+public class RequestOptions {
+  private Function<Request, Request> authenticateFunc;
+  private Function<Request, Request> urlFunc;
+  private Function<Request, Request> userAgentFunc;
+
+  /**
+   * Constructs a new RequestOptions instance with default identity functions. Initially, all
+   * transformations are set to pass through the request unchanged.
+   */
+  public RequestOptions() {
+    // Default to identity functions
+    this.authenticateFunc = request -> request;
+    this.urlFunc = request -> request;
+    this.userAgentFunc = request -> request;
+  }
+
+  /**
+   * Sets the authorization header for the request.
+   *
+   * @param authorization The authorization value to be set in the header
+   * @return This RequestOptions instance for method chaining
+   */
+  public RequestOptions withAuthorization(String authorization) {
+    this.authenticateFunc = request -> request.withHeader("Authorization", authorization);
+    return this;
+  }
+
+  /**
+   * Sets the URL for the request.
+   *
+   * @param url The URL to be set for the request
+   * @return This RequestOptions instance for method chaining
+   */
+  public RequestOptions withUrl(String url) {
+    this.urlFunc = request -> request.withUrl(url);
+    return this;
+  }
+
+  /**
+   * Sets the User-Agent header for the request.
+   *
+   * @param userAgent The user agent string to be set in the header
+   * @return This RequestOptions instance for method chaining
+   */
+  public RequestOptions withUserAgent(String userAgent) {
+    this.userAgentFunc = request -> request.withHeader("User-Agent", userAgent);
+    return this;
+  }
+
+  /**
+   * Applies all configured transformations to the given request. The transformations are applied in
+   * the following order: 1. Authentication 2. URL 3. User-Agent
+   *
+   * @param request The original request to be transformed
+   * @return A new Request instance with all transformations applied
+   */
+  public Request applyOptions(Request request) {
+    // Apply all transformation functions in sequence
+    return userAgentFunc.apply(urlFunc.apply(authenticateFunc.apply(request)));
+  }
+}