diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index f881abc99..ee44d5cbc 100644 --- a/NEXT_CHANGELOG.md +++ b/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 diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java index 1e75e021d..b1e9ca61a 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/commons/CommonsHttpClient.java +++ b/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,7 +66,9 @@ 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) { @@ -73,6 +76,17 @@ public Builder withTimeoutSeconds(int 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) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/commons/CommonsHttpClientTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/commons/CommonsHttpClientTest.java index 0e9df2be6..b8a45ecde 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/commons/CommonsHttpClientTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/commons/CommonsHttpClientTest.java @@ -7,9 +7,12 @@ import com.databricks.sdk.core.http.Request; import com.databricks.sdk.core.http.Response; import java.io.IOException; +import java.lang.reflect.Field; import java.net.HttpURLConnection; import java.util.*; import org.apache.commons.io.IOUtils; +import org.apache.http.client.config.RequestConfig; +import org.apache.http.client.methods.Configurable; import org.junit.jupiter.api.Test; class CommonsHttpClientTest { @@ -88,6 +91,24 @@ public void testNoRedirection() throws IOException { } } + @Test + public void testCustomRequestConfig() throws Exception { + RequestConfig requestConfig = + RequestConfig.custom() + .setConnectTimeout(10_000) + .setSocketTimeout(900_000) + .setConnectionRequestTimeout(300_000) + .build(); + CommonsHttpClient httpClient = + new CommonsHttpClient.Builder().withRequestConfig(requestConfig).build(); + + // Verify that the custom RequestConfig is passed to the underlying Apache HttpClient. + Field hcField = CommonsHttpClient.class.getDeclaredField("hc"); + hcField.setAccessible(true); + Configurable internalClient = (Configurable) hcField.get(httpClient); + assertSame(requestConfig, internalClient.getConfig()); + } + @Test public void testRedirection() throws IOException { List fixtures =