Merge pull request #176 from agnipratimnag-db/agnipratim_peco-1642

[PECO-1642] Added configurability of Polling Interval

Author: agnipratimnag-db

src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksSdkClient.java

@@ -37,7 +37,6 @@ public class DatabricksSdkClient implements DatabricksClient {
   private static final Logger LOGGER = LoggerFactory.getLogger(DatabricksSdkClient.class);
   private static final String SYNC_TIMEOUT_VALUE = "10s";
   private static final String ASYNC_TIMEOUT_VALUE = "0s";
-  private static final int STATEMENT_RESULT_POLL_INTERVAL_MILLIS = 200;
 
   private final IDatabricksConnectionContext connectionContext;
   private final DatabricksConfig databricksConfig;
@@ -160,7 +159,7 @@ public DatabricksResultSet executeStatement(
     while (responseState == StatementState.PENDING || responseState == StatementState.RUNNING) {
       if (pollCount > 0) { // First poll happens without a delay
         try {
-          Thread.sleep(STATEMENT_RESULT_POLL_INTERVAL_MILLIS); // TODO: make this configurable
+          Thread.sleep(this.connectionContext.getAsyncExecPollInterval());
         } catch (InterruptedException e) {
           throw new DatabricksTimeoutException("Thread interrupted due to statement timeout");
         }

src/main/java/com/databricks/jdbc/driver/DatabricksConnectionContext.java

@@ -172,6 +172,13 @@ public String getToken() {
         : getParameter(DatabricksJdbcConstants.PWD);
   }
 
+  @Override
+  public int getAsyncExecPollInterval() {
+    return getParameter(POLL_INTERVAL) == null
+        ? POLL_INTERVAL_DEFAULT
+        : Integer.parseInt(getParameter(DatabricksJdbcConstants.POLL_INTERVAL));
+  }
+
   public String getCloud() throws DatabricksParsingException {
     String hostURL = getHostUrl();
     if (hostURL.contains("azuredatabricks.net")

src/main/java/com/databricks/jdbc/driver/DatabricksJdbcConstants.java

@@ -58,6 +58,9 @@ public final class DatabricksJdbcConstants {
   static final String UID = "uid";
   static final String PWD = "pwd";
 
+  static final String POLL_INTERVAL = "asyncexecpollinterval";
+  static final int POLL_INTERVAL_DEFAULT = 200;
+
   static final String AWS_CLIENT_ID = "databricks-sql-jdbc";
 
   static final String AAD_CLIENT_ID = "96eecda7-19ea-49cc-abb5-240097d554f5";

src/main/java/com/databricks/jdbc/driver/IDatabricksConnectionContext.java

@@ -114,5 +114,7 @@ public static AuthMech parseAuthMech(String authMech) {
 
   String getEndpointURL() throws DatabricksParsingException;
 
+  int getAsyncExecPollInterval();
+
   DatabricksClientType getClientType();
 }

src/test/java/com/databricks/jdbc/driver/DatabricksConnectionContextTest.java

@@ -52,6 +52,10 @@ class DatabricksConnectionContextTest {
       "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/5c89f447c476a5a8;UseSystemProxy=1;UseProxy=1;ProxyHost=127.0.0.1;ProxyPort=8080;ProxyAuth=1;ProxyUID=proxyUser;ProxyPwd=proxyPassword;UseCFProxy=1;CFProxyHost=127.0.1.2;CFProxyPort=8081;CFProxyAuth=1;CFProxyUID=cfProxyUser;CFProxyPwd=cfProxyPassword;";
 
   private static Properties properties = new Properties();
+
+  private static final String VALID_URL_POLLING =
+      "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:4473;ssl=1;asyncexecpollinterval=500;AuthMech=3;httpPath=/sql/1.0/warehouses/5c89f447c476a5a8;QueryResultCompressionType=1";
+
   private static Properties properties_with_pwd = new Properties();
 
   @BeforeAll
@@ -246,6 +250,18 @@ public void testParsingOfUrlWithoutDefault() throws DatabricksSQLException {
     assertEquals("INFO", connectionContext.getLogLevelString());
   }
 
+  @Test
+  public void testPollingInterval() throws DatabricksSQLException {
+    DatabricksConnectionContext connectionContext =
+        (DatabricksConnectionContext) DatabricksConnectionContext.parse(VALID_URL_5, properties);
+    assertEquals(200, connectionContext.getAsyncExecPollInterval());
+
+    DatabricksConnectionContext connectionContextWithPoll =
+        (DatabricksConnectionContext)
+            DatabricksConnectionContext.parse(VALID_URL_POLLING, properties);
+    assertEquals(500, connectionContextWithPoll.getAsyncExecPollInterval());
+  }
+
   @Test
   public void testParsingOfUrlWithProxy() throws DatabricksSQLException {
     IDatabricksConnectionContext connectionContext =