@@ -47,6 +47,7 @@ public class CommonsHttpClient implements HttpClient {
4747 public static class Builder {
4848 private DatabricksConfig databricksConfig ;
4949 private Integer timeoutSeconds ;
50+ private RequestConfig requestConfig ;
5051 private ProxyConfig proxyConfig ;
5152 private SSLConnectionSocketFactory sslSocketFactory ;
5253 private PoolingHttpClientConnectionManager connectionManager ;
@@ -65,14 +66,27 @@ public Builder withDatabricksConfig(DatabricksConfig databricksConfig) {
6566
6667 /**
6768 * @param timeoutSeconds The timeout in seconds to use for the HttpClient. This will override
68- * any timeout set in the DatabricksConfig.
69+ * any timeout set in the DatabricksConfig. Sets all three Apache HttpClient timeouts
70+ * (connect, socket, connection request) to the same value. For fine-grained control, use
71+ * {@link #withRequestConfig(RequestConfig)} instead.
6972 * @return This builder.
7073 */
7174 public Builder withTimeoutSeconds (int timeoutSeconds ) {
7275 this .timeoutSeconds = timeoutSeconds ;
7376 return this ;
7477 }
7578
79+ /**
80+ * @param requestConfig The Apache {@link RequestConfig} to use for the HttpClient. When set,
81+ * this takes precedence over {@link #withTimeoutSeconds(int)} and any timeout from {@link
82+ * DatabricksConfig}.
83+ * @return This builder.
84+ */
85+ public Builder withRequestConfig (RequestConfig requestConfig ) {
86+ this .requestConfig = requestConfig ;
87+ return this ;
88+ }
89+
7690 /**
7791 * @param proxyConfig the proxy configuration to use for the HttpClient.
7892 * @return This builder.
@@ -121,17 +135,12 @@ public CommonsHttpClient build() {
121135 private final CloseableHttpClient hc ;
122136
123137 private CommonsHttpClient (Builder builder ) {
124- int timeoutSeconds = 300 ;
125- if (builder .databricksConfig != null
126- && builder .databricksConfig .getHttpTimeoutSeconds () != null ) {
127- timeoutSeconds = builder .databricksConfig .getHttpTimeoutSeconds ();
128- }
129- if (builder .timeoutSeconds != null ) {
130- timeoutSeconds = builder .timeoutSeconds ;
131- }
132- int timeout = timeoutSeconds * 1000 ;
138+ RequestConfig requestConfig =
139+ builder .requestConfig != null
140+ ? builder .requestConfig
141+ : makeDefaultRequestConfig (builder .databricksConfig , builder .timeoutSeconds );
133142 HttpClientBuilder httpClientBuilder =
134- HttpClientBuilder .create ().setDefaultRequestConfig (makeRequestConfig ( timeout ) );
143+ HttpClientBuilder .create ().setDefaultRequestConfig (requestConfig );
135144 if (builder .proxyConfig != null ) {
136145 ProxyUtils .setupProxy (builder .proxyConfig , httpClientBuilder );
137146 }
@@ -156,7 +165,16 @@ private CommonsHttpClient(Builder builder) {
156165 hc = httpClientBuilder .build ();
157166 }
158167
159- private RequestConfig makeRequestConfig (int timeout ) {
168+ private static RequestConfig makeDefaultRequestConfig (
169+ DatabricksConfig databricksConfig , Integer timeoutSecondsOverride ) {
170+ int timeoutSeconds = 300 ;
171+ if (databricksConfig != null && databricksConfig .getHttpTimeoutSeconds () != null ) {
172+ timeoutSeconds = databricksConfig .getHttpTimeoutSeconds ();
173+ }
174+ if (timeoutSecondsOverride != null ) {
175+ timeoutSeconds = timeoutSecondsOverride ;
176+ }
177+ int timeout = timeoutSeconds * 1000 ;
160178 return RequestConfig .custom ()
161179 .setConnectionRequestTimeout (timeout )
162180 .setConnectTimeout (timeout )
0 commit comments