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 0af1d2231..fb4817326 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 @@ -23,6 +23,7 @@ import java.util.stream.Collectors; import org.apache.commons.io.IOUtils; import org.apache.http.*; +import org.apache.http.client.HttpRequestRetryHandler; import org.apache.http.client.config.RequestConfig; import org.apache.http.client.methods.*; import org.apache.http.conn.ssl.SSLConnectionSocketFactory; @@ -47,6 +48,7 @@ public static class Builder { private ProxyConfig proxyConfig; private SSLConnectionSocketFactory sslSocketFactory; private PoolingHttpClientConnectionManager connectionManager; + private HttpRequestRetryHandler requestRetryHandler; /** * @param databricksConfig The DatabricksConfig to use for the HttpClient. If the @@ -96,6 +98,17 @@ public Builder withConnectionManager(PoolingHttpClientConnectionManager connecti return this; } + /** + * @param requestRetryHandler the HttpRequestRetryHandler to use for the HttpClient. + * @return This builder. + *

Note: This API is experimental and may change or be removed in future releases + * without notice. + */ + public Builder withRequestRetryHandler(HttpRequestRetryHandler requestRetryHandler) { + this.requestRetryHandler = requestRetryHandler; + return this; + } + /** Builds a new instance of CommonsHttpClient with the configured parameters. */ public CommonsHttpClient build() { return new CommonsHttpClient(this); @@ -131,6 +144,9 @@ private CommonsHttpClient(Builder builder) { connectionManager.setMaxTotal(100); httpClientBuilder.setConnectionManager(connectionManager); } + if (builder.requestRetryHandler != null) { + httpClientBuilder.setRetryHandler(builder.requestRetryHandler); + } hc = httpClientBuilder.build(); }