Solved PECO-1642

Author: agnipratimnag-db

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

@@ -37,7 +37,7 @@ 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 int STATEMENT_RESULT_POLL_INTERVAL_MILLIS = 200;
 
   private final IDatabricksConnectionContext connectionContext;
   private final DatabricksConfig databricksConfig;
@@ -61,6 +61,7 @@ public DatabricksSdkClient(IDatabricksConnectionContext connectionContext)
 
     OAuthAuthenticator authenticator = new OAuthAuthenticator(connectionContext);
     this.workspaceClient = authenticator.getWorkspaceClient();
+    this.STATEMENT_RESULT_POLL_INTERVAL_MILLIS = connectionContext.getInterval();
   }
 
   public DatabricksSdkClient(
@@ -77,6 +78,8 @@ public DatabricksSdkClient(
     this.workspaceClient =
         new WorkspaceClient(true /* mock */, apiClient)
             .withStatementExecutionImpl(statementExecutionService);
+
+    this.STATEMENT_RESULT_POLL_INTERVAL_MILLIS = connectionContext.getInterval();
   }
 
   @Override
@@ -160,7 +163,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(STATEMENT_RESULT_POLL_INTERVAL_MILLIS);
         } 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 getInterval() {
+    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

@@ -57,6 +57,8 @@ 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 getInterval();
+
   DatabricksClientType getClientType();
 }

src/test/java/com/databricks/jdbc/local/DriverTester.java

@@ -36,7 +36,7 @@ void testGetTablesOSS_StatementExecution() throws Exception {
     DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
     // Getting the connection
     String jdbcUrl =
-        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
+        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;asyncexecpollinterval=200;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
     Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "x");
     System.out.println("Connection established......");
     Statement statement = con.createStatement();
@@ -54,11 +54,13 @@ void testGetTablesOSS_Metadata() throws Exception {
     DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
     // Getting the connection
     String jdbcUrl =
-        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
-    Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "x");
+        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;asyncexecpollinterval=200;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
+    Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "dapiecd97706ac974a03b199ec117cf0bdf6");
     System.out.println("Connection established......");
-    DatabaseMetaData metaData = con.getMetaData();
-    ResultSet resultSet = metaData.getCatalogs();
+//    DatabaseMetaData metaData = con.getMetaData();
+//    ResultSet resultSet = metaData.getCatalogs();
+    Statement statement = con.createStatement();
+    ResultSet resultSet = statement.executeQuery("select * from main.newschema_agni.agni_table");
     printResultSet(resultSet);
     resultSet.close();
     con.close();
@@ -70,7 +72,7 @@ void testArclight() throws Exception {
     DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
     // Getting the connection
     String jdbcUrl =
-        "jdbc:databricks://arclight-staging-e2-arclight-dmk-qa-staging-us-east-1.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/8561171c1d9afb1f;";
+        "jdbc:databricks://arclight-staging-e2-arclight-dmk-qa-staging-us-east-1.staging.cloud.databricks.com:443/default;transportMode=https;asyncexecpollinterval=200;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/8561171c1d9afb1f;";
     Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "xx");
     System.out.println("Connection established......");
     // Retrieving data
@@ -91,7 +93,7 @@ void testAllPurposeClusters() throws Exception {
     DriverManager.registerDriver(new com.databricks.jdbc.driver.DatabricksDriver());
     DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
     String jdbcUrl =
-        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UID=token;";
+        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;asyncexecpollinterval=200;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UID=token;";
     Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "x");
     System.out.println("Connection established......");
     Statement statement = con.createStatement();
@@ -107,7 +109,7 @@ void testAllPurposeClustersMetadata() throws Exception {
     DriverManager.registerDriver(new com.databricks.jdbc.driver.DatabricksDriver());
     DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
     String jdbcUrl =
-        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UID=token;";
+        "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;asyncexecpollinterval=200;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UID=token;";
     Connection con = DriverManager.getConnection(jdbcUrl, "[email protected]", "xx");
     System.out.println("Connection established......");
     ResultSet resultSet = con.getMetaData().getPrimaryKeys("main", "ggm_pk", "table_with_pk");