[PECOBLR-164] Support connection param SocketTimeout (#777)

- The default value of the param is kept as 900 seconds (same as earlier value and default in databricks-sdk-java)
- Set the param in thrift, dbfs, and sql-exec clients

Author: jayantsing-db

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java

@@ -44,8 +44,8 @@ public class DatabricksConnection implements IDatabricksConnection, IDatabricksC
   public DatabricksConnection(IDatabricksConnectionContext connectionContext)
       throws DatabricksSQLException {
     this.connectionContext = connectionContext;
-    this.session = new DatabricksSession(connectionContext);
     DatabricksThreadContextHolder.setConnectionContext(connectionContext);
+    this.session = new DatabricksSession(connectionContext);
   }
 
   @VisibleForTesting

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

@@ -824,6 +824,12 @@ public int getRowsFetchedPerBlock() {
     return maxRows;
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public int getSocketTimeout() {
+    return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.SOCKET_TIMEOUT));
+  }
+
   private static boolean nullOrEmptyString(String s) {
     return s == null || s.isEmpty();
   }

src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java

@@ -257,4 +257,7 @@ public interface IDatabricksConnectionContext {
 
   /** Returns maximum number of rows that a query returns at a time. */
   int getRowsFetchedPerBlock();
+
+  /** Returns the socket timeout in seconds for HTTP connections. */
+  int getSocketTimeout();
 }

src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java

@@ -108,7 +108,8 @@ public enum DatabricksJdbcUrlParams {
   DEFAULT_STRING_COLUMN_LENGTH(
       "DefaultStringColumnLength",
       "Maximum number of characters that can be contained in STRING columns",
-      "255");
+      "255"),
+  SOCKET_TIMEOUT("socketTimeout", "Socket timeout in seconds", "900");
 
   private final String paramName;
   private final String defaultValue;

src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java

@@ -43,6 +43,7 @@ public ClientConfigurator(IDatabricksConnectionContext connectionContext) {
     this.connectionContext = connectionContext;
     this.databricksConfig = new DatabricksConfig();
     CommonsHttpClient.Builder httpClientBuilder = new CommonsHttpClient.Builder();
+    httpClientBuilder.withTimeoutSeconds(connectionContext.getSocketTimeout());
     setupProxyConfig(httpClientBuilder);
     setupConnectionManager(httpClientBuilder);
     this.databricksConfig.setHttpClient(httpClientBuilder.build());

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

@@ -45,8 +45,6 @@ public class DatabricksHttpClient implements IDatabricksHttpClient, Closeable {
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksHttpClient.class);
   private static final int DEFAULT_MAX_HTTP_CONNECTIONS = 1000;
   private static final int DEFAULT_MAX_HTTP_CONNECTIONS_PER_ROUTE = 1000;
-  private static final int DEFAULT_HTTP_CONNECTION_TIMEOUT = 60 * 1000; // ms
-  private static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300 * 1000; // ms
   private final PoolingHttpClientConnectionManager connectionManager;
   private final CloseableHttpClient httpClient;
   private IdleConnectionEvictor idleConnectionEvictor;
@@ -134,11 +132,12 @@ private PoolingHttpClientConnectionManager initializeConnectionManager(
     return connectionManager;
   }
 
-  private RequestConfig makeRequestConfig() {
+  private RequestConfig makeRequestConfig(int timeoutSeconds) {
+    int timeoutMillis = timeoutSeconds * 1000;
     return RequestConfig.custom()
-        .setConnectionRequestTimeout(DEFAULT_HTTP_CONNECTION_TIMEOUT)
-        .setConnectTimeout(DEFAULT_HTTP_CONNECTION_TIMEOUT)
-        .setSocketTimeout(DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT)
+        .setConnectionRequestTimeout(timeoutMillis)
+        .setConnectTimeout(timeoutMillis)
+        .setSocketTimeout(timeoutMillis)
         .build();
   }
 
@@ -152,7 +151,7 @@ private CloseableHttpClient makeClosableHttpClient(
         HttpClientBuilder.create()
             .setConnectionManager(connectionManager)
             .setUserAgent(UserAgentManager.getUserAgentString())
-            .setDefaultRequestConfig(makeRequestConfig())
+            .setDefaultRequestConfig(makeRequestConfig(connectionContext.getSocketTimeout()))
             .setRetryHandler(retryHandler)
             .addInterceptorFirst(retryHandler);
     setupProxy(connectionContext, builder);