Add withRequestConfig to CommonsHttpClient.Builder

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.97.0
 
 ### New Features and Improvements
+ * Add `withRequestConfig(RequestConfig)` to `CommonsHttpClient.Builder` for fine-grained timeout control ([#691](https://github.com/databricks/databricks-sdk-java/pull/691)).
 
 ### Bug Fixes
 

databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java

@@ -47,6 +47,7 @@ public class CommonsHttpClient implements HttpClient {
   public static class Builder {
     private DatabricksConfig databricksConfig;
     private Integer timeoutSeconds;
+    private RequestConfig requestConfig;
     private ProxyConfig proxyConfig;
     private SSLConnectionSocketFactory sslSocketFactory;
     private PoolingHttpClientConnectionManager connectionManager;
@@ -65,14 +66,27 @@ public Builder withDatabricksConfig(DatabricksConfig databricksConfig) {
 
     /**
      * @param timeoutSeconds The timeout in seconds to use for the HttpClient. This will override
-     *     any timeout set in the DatabricksConfig.
+     *     any timeout set in the DatabricksConfig. Sets all three Apache HttpClient timeouts
+     *     (connect, socket, connection request) to the same value. For fine-grained control, use
+     *     {@link #withRequestConfig(RequestConfig)} instead.
      * @return This builder.
      */
     public Builder withTimeoutSeconds(int timeoutSeconds) {
       this.timeoutSeconds = timeoutSeconds;
       return this;
     }
 
+    /**
+     * @param requestConfig The Apache {@link RequestConfig} to use for the HttpClient. When set,
+     *     this takes precedence over {@link #withTimeoutSeconds(int)} and any timeout from {@link
+     *     DatabricksConfig}.
+     * @return This builder.
+     */
+    public Builder withRequestConfig(RequestConfig requestConfig) {
+      this.requestConfig = requestConfig;
+      return this;
+    }
+
     /**
      * @param proxyConfig the proxy configuration to use for the HttpClient.
      * @return This builder.
@@ -121,17 +135,12 @@ public CommonsHttpClient build() {
   private final CloseableHttpClient hc;
 
   private CommonsHttpClient(Builder builder) {
-    int timeoutSeconds = 300;
-    if (builder.databricksConfig != null
-        && builder.databricksConfig.getHttpTimeoutSeconds() != null) {
-      timeoutSeconds = builder.databricksConfig.getHttpTimeoutSeconds();
-    }
-    if (builder.timeoutSeconds != null) {
-      timeoutSeconds = builder.timeoutSeconds;
-    }
-    int timeout = timeoutSeconds * 1000;
+    RequestConfig requestConfig =
+        builder.requestConfig != null
+            ? builder.requestConfig
+            : makeDefaultRequestConfig(builder.databricksConfig, builder.timeoutSeconds);
     HttpClientBuilder httpClientBuilder =
-        HttpClientBuilder.create().setDefaultRequestConfig(makeRequestConfig(timeout));
+        HttpClientBuilder.create().setDefaultRequestConfig(requestConfig);
     if (builder.proxyConfig != null) {
       ProxyUtils.setupProxy(builder.proxyConfig, httpClientBuilder);
     }
@@ -156,7 +165,16 @@ private CommonsHttpClient(Builder builder) {
     hc = httpClientBuilder.build();
   }
 
-  private RequestConfig makeRequestConfig(int timeout) {
+  private static RequestConfig makeDefaultRequestConfig(
+      DatabricksConfig databricksConfig, Integer timeoutSecondsOverride) {
+    int timeoutSeconds = 300;
+    if (databricksConfig != null && databricksConfig.getHttpTimeoutSeconds() != null) {
+      timeoutSeconds = databricksConfig.getHttpTimeoutSeconds();
+    }
+    if (timeoutSecondsOverride != null) {
+      timeoutSeconds = timeoutSecondsOverride;
+    }
+    int timeout = timeoutSeconds * 1000;
     return RequestConfig.custom()
         .setConnectionRequestTimeout(timeout)
         .setConnectTimeout(timeout)

databricks-sdk-java/src/test/java/com/databricks/sdk/core/commons/CommonsHttpClientTest.java

@@ -10,6 +10,7 @@
 import java.net.HttpURLConnection;
 import java.util.*;
 import org.apache.commons.io.IOUtils;
+import org.apache.http.client.config.RequestConfig;
 import org.junit.jupiter.api.Test;
 
 class CommonsHttpClientTest {
@@ -88,6 +89,23 @@ public void testNoRedirection() throws IOException {
     }
   }
 
+  @Test
+  public void testCustomRequestConfig() throws IOException {
+    try (FixtureServer fixtures = new FixtureServer().with("GET", "/foo", "ok", 200)) {
+      RequestConfig requestConfig =
+          RequestConfig.custom()
+              .setConnectTimeout(10_000)
+              .setSocketTimeout(900_000)
+              .setConnectionRequestTimeout(300_000)
+              .build();
+      HttpClient httpClient =
+          new CommonsHttpClient.Builder().withRequestConfig(requestConfig).build();
+      Request in = new Request("GET", fixtures.getUrl() + "/foo");
+      Response out = httpClient.execute(in);
+      assertEquals("ok", out.getDebugBody().trim());
+    }
+  }
+
   @Test
   public void testRedirection() throws IOException {
     List<FixtureServer.FixtureMapping> fixtures =