Run queries in async mode in thrift client (#570)

- Run execution queries in `async` mode
- Refactor result retrieval in Thrift by improving response handling and fixing polling
- Competing driver, SQL Gateway, and SQL Execution API also execute the queries in async mode
- Minor code cleanup

Running queries in async mode and using polling will be beneficial:
- Freeing up server resources under high load instead of holding the connection for ~5 seconds (direct results)
- If query takes more than 5 seconds, server will anyway require polling from client
- Didn't notice any regressions (slightly better performance) under high concurrent load when running queries in async mode. Some empirical results below:
  - Running a simple select query on a table with 10-15 columns for 10K rows in a concurrent environment (250 connections executing query in parallel)
- Average results when running on `main`
```
Average query time across connections: 73873.58ms
Min query time: 22075ms
Max query time: 88997ms
```
- Average results when running in `async` mode
```
Average query time across connections: 68382.95ms
Min query time: 24718ms
Max query time: 95282ms
```

Author: jayantsing-db

.gitignore

@@ -6,3 +6,4 @@ TestFile
 logs/
 *.log
 *.pem
+dependency-reduced-pom.xml

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowResultChunk.java

@@ -1,11 +1,11 @@
 package com.databricks.jdbc.api.impl.arrow;
 
-import static com.databricks.jdbc.common.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
 import static com.databricks.jdbc.common.util.DatabricksThriftUtil.createExternalLink;
 import static com.databricks.jdbc.common.util.ValidationUtil.checkHTTPError;
 
 import com.databricks.jdbc.common.CompressionCodec;
 import com.databricks.jdbc.common.util.DecompressionUtil;
+import com.databricks.jdbc.common.util.DriverUtil;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
@@ -215,7 +215,7 @@ void setStatus(ChunkStatus status) {
   /** Checks if the link is valid */
   boolean isChunkLinkInvalid() {
     return status == ChunkStatus.PENDING
-        || (!Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))
+        || (!DriverUtil.isRunningAgainstFake()
             && expiryTime.minusSeconds(SECONDS_BUFFER_FOR_EXPIRY).isBefore(Instant.now()));
   }
 

src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java

@@ -22,7 +22,6 @@
 public class DatabricksThriftUtil {
 
   private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksThriftUtil.class);
-
   public static final List<TStatusCode> SUCCESS_STATUS_LIST =
       List.of(TStatusCode.SUCCESS_STATUS, TStatusCode.SUCCESS_WITH_INFO_STATUS);
 
@@ -317,15 +316,15 @@ public static void checkDirectResultsForErrorStatus(
     }
     if (directResults.isSetResultSetMetadata()) {
       LOGGER.debug("direct results metadata being verified for success response");
-      verifySuccessStatus(directResults.getResultSetMetadata().status, context);
+      verifySuccessStatus(directResults.getResultSetMetadata().getStatus(), context);
     }
     if (directResults.isSetCloseOperation()) {
       LOGGER.debug("direct results close operation verified for success response");
-      verifySuccessStatus(directResults.getCloseOperation().status, context);
+      verifySuccessStatus(directResults.getCloseOperation().getStatus(), context);
     }
     if (directResults.isSetResultSet()) {
       LOGGER.debug("direct result set being verified for success response");
-      verifySuccessStatus(directResults.getResultSet().status, context);
+      verifySuccessStatus(directResults.getResultSet().getStatus(), context);
     }
   }
 }

src/main/java/com/databricks/jdbc/common/util/DriverUtil.java

@@ -26,7 +26,6 @@ public class DriverUtil {
   private static final String DBSQL_VERSION_SQL = "SELECT current_version().dbsql_version";
   public static final int DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT = 2024;
   public static final int DBSQL_MIN_MINOR_VERSION_FOR_SEA_SUPPORT = 30;
-
   private static final String[] VERSION_PARTS = VERSION.split("[.-]");
 
   public static String getVersion() {
@@ -86,7 +85,7 @@ private static Throwable getRootCause(Throwable throwable) {
   static void ensureUpdatedDBRVersionInUse(IDatabricksConnection connection)
       throws DatabricksValidationException {
     if (connection.getConnectionContext().getClientType() != DatabricksClientType.SQL_EXEC
-        || Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
+        || isRunningAgainstFake()) {
       // Check applicable only for SEA flow
       return;
     }
@@ -127,4 +126,8 @@ private static boolean doesDriverSupportSEA(String dbsqlVersion) {
     }
     return majorVersion > DBSQL_MIN_MAJOR_VERSION_FOR_SEA_SUPPORT;
   }
+
+  public static boolean isRunningAgainstFake() {
+    return Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP));
+  }
 }

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

@@ -5,6 +5,7 @@
 import static io.netty.util.NetUtil.LOCALHOST;
 
 import com.databricks.jdbc.api.IDatabricksConnectionContext;
+import com.databricks.jdbc.common.util.DriverUtil;
 import com.databricks.jdbc.common.util.UserAgentManager;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.dbclient.impl.common.ConfiguratorUtils;
@@ -66,7 +67,7 @@ public class DatabricksHttpClient implements IDatabricksHttpClient, Closeable {
   public CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHttpException {
     LOGGER.debug(
         String.format("Executing HTTP request [{%s}]", RequestSanitizer.sanitizeRequest(request)));
-    if (!Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
+    if (!DriverUtil.isRunningAgainstFake()) {
       // TODO : allow gzip in wiremock
       request.setHeader("Content-Encoding", "gzip");
     }
@@ -118,7 +119,7 @@ private CloseableHttpClient makeClosableHttpClient(
             .setRetryHandler(retryHandler)
             .addInterceptorFirst(retryHandler);
     setupProxy(connectionContext, builder);
-    if (Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
+    if (DriverUtil.isRunningAgainstFake()) {
       setFakeServiceRouteInHttpClient(builder);
     }
     return builder.build();

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java

@@ -1,15 +1,14 @@
 package com.databricks.jdbc.dbclient.impl.thrift;
 
-import static com.databricks.jdbc.common.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
 import static com.databricks.jdbc.common.EnvironmentVariables.*;
 import static com.databricks.jdbc.common.util.DatabricksThriftUtil.*;
-import static com.databricks.jdbc.model.client.thrift.generated.TStatusCode.*;
 
 import com.databricks.jdbc.api.IDatabricksConnectionContext;
 import com.databricks.jdbc.api.IDatabricksSession;
 import com.databricks.jdbc.api.impl.*;
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.common.StatementType;
+import com.databricks.jdbc.common.util.DriverUtil;
 import com.databricks.jdbc.dbclient.impl.common.ClientConfigurator;
 import com.databricks.jdbc.dbclient.impl.common.StatementId;
 import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory;
@@ -47,9 +46,7 @@ final class DatabricksThriftAccessor {
     Map<String, String> authHeaders = databricksConfig.authenticate();
     String endPointUrl = connectionContext.getEndpointURL();
 
-    final boolean isFakeServiceTest =
-        Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP));
-    if (!isFakeServiceTest) {
+    if (!DriverUtil.isRunningAgainstFake()) {
       // Create a new thrift client for each thread as client state is not thread safe. Note that
       // the underlying protocol uses the same http client which is thread safe
       this.thriftClient =
@@ -127,7 +124,7 @@ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, String cont
       throws DatabricksHttpException {
     refreshHeadersIfRequired();
     return getResultSetResp(
-        new TStatus().setStatusCode(SUCCESS_STATUS),
+        new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS),
         operationHandle,
         context,
         DEFAULT_ROW_LIMIT,
@@ -170,7 +167,7 @@ TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement)
             parentStatement.getStatementId().toSQLExecStatementId());
     int maxRows = (parentStatement == null) ? DEFAULT_ROW_LIMIT : parentStatement.getMaxRows();
     return getResultSetResp(
-        new TStatus().setStatusCode(SUCCESS_STATUS),
+        new TStatus().setStatusCode(TStatusCode.SUCCESS_STATUS),
         getOperationHandle(parentStatement.getStatementId()),
         context,
         maxRows,
@@ -213,50 +210,53 @@ private TFetchResultsResp getResultSetResp(
     return response;
   }
 
-  private void longPolling(TOperationHandle operationHandle)
-      throws TException, InterruptedException, DatabricksHttpException {
-    TGetOperationStatusReq request =
-        new TGetOperationStatusReq()
-            .setOperationHandle(operationHandle)
-            .setGetProgressUpdate(false);
-    TGetOperationStatusResp response;
-    TStatus status;
-    do {
-      response = getThriftClient().GetOperationStatus(request);
-      status = response.getStatus();
-      if (status.getStatusCode() == TStatusCode.STILL_EXECUTING_STATUS) {
-        Thread.sleep(DEFAULT_SLEEP_DELAY);
-      }
-    } while (status.getStatusCode() == TStatusCode.STILL_EXECUTING_STATUS);
-    verifySuccessStatus(status, String.format("Request {%s}, Response {%s}", request, response));
-  }
-
   DatabricksResultSet execute(
       TExecuteStatementReq request,
       IDatabricksStatementInternal parentStatement,
       IDatabricksSession session,
       StatementType statementType)
       throws SQLException {
     refreshHeadersIfRequired();
+
+    // Set direct result configuration
     int maxRows = (parentStatement == null) ? DEFAULT_ROW_LIMIT : parentStatement.getMaxRows();
     if (enableDirectResults) {
       TSparkGetDirectResults directResults =
           new TSparkGetDirectResults().setMaxBytes(DEFAULT_BYTE_LIMIT).setMaxRows(maxRows);
       request.setGetDirectResults(directResults);
     }
+
     TExecuteStatementResp response;
-    TFetchResultsResp resultSet = null;
+    TFetchResultsResp resultSet;
+
     try {
       response = getThriftClient().ExecuteStatement(request);
-      if (Arrays.asList(ERROR_STATUS, INVALID_HANDLE_STATUS).contains(response.status.statusCode)) {
-        throw new DatabricksSQLException(response.status.errorMessage, response.status.sqlState);
-      }
+      checkResponseForErrors(response);
+
+      TGetOperationStatusResp statusResp = null;
       if (response.isSetDirectResults()) {
         checkDirectResultsForErrorStatus(response.getDirectResults(), response.toString());
+        statusResp = response.getDirectResults().getOperationStatus();
+        checkOperationStatusForErrors(statusResp);
+      }
+
+      // Polling until query operation state is finished
+      TGetOperationStatusReq statusReq =
+          new TGetOperationStatusReq()
+              .setOperationHandle(response.getOperationHandle())
+              .setGetProgressUpdate(false);
+      while (shouldContinuePolling(statusResp)) {
+        statusResp = getThriftClient().GetOperationStatus(statusReq);
+        checkOperationStatusForErrors(statusResp);
+      }
+
+      if (hasResultDataInDirectResults(response)) {
+        // The first response has result data
+        // There is no polling in this case as status was already finished
         resultSet = response.getDirectResults().getResultSet();
         resultSet.setResultSetMetadata(response.getDirectResults().getResultSetMetadata());
       } else {
-        longPolling(response.getOperationHandle());
+        // Fetch the result data after polling
         resultSet =
             getResultSetResp(
                 response.getStatus(),
@@ -265,11 +265,11 @@ DatabricksResultSet execute(
                 maxRows,
                 true);
       }
-    } catch (TException | InterruptedException e) {
+    } catch (TException e) {
       String errorMessage =
           String.format(
               "Error while receiving response from Thrift server. Request {%s}, Error {%s}",
-              request.toString(), e.getMessage());
+              request, e.getMessage());
       LOGGER.error(e, errorMessage);
       throw new DatabricksHttpException(errorMessage, e);
     }
@@ -288,12 +288,11 @@ DatabricksResultSet executeAsync(
       StatementType statementType)
       throws SQLException {
     refreshHeadersIfRequired();
-    TExecuteStatementResp response = null;
-    DatabricksHttpTTransport transport =
-        (DatabricksHttpTTransport) getThriftClient().getInputProtocol().getTransport();
+    TExecuteStatementResp response;
     try {
       response = getThriftClient().ExecuteStatement(request);
-      if (Arrays.asList(ERROR_STATUS, INVALID_HANDLE_STATUS).contains(response.status.statusCode)) {
+      if (Arrays.asList(TStatusCode.ERROR_STATUS, TStatusCode.INVALID_HANDLE_STATUS)
+          .contains(response.status.statusCode)) {
         LOGGER.error(
             "Received error response {%s} from Thrift Server for request {%s}",
             response, request.toString());
@@ -335,7 +334,8 @@ DatabricksResultSet getStatementResult(
     try {
       response = getThriftClient().GetOperationStatus(request);
       statusCode = response.getStatus().getStatusCode();
-      if (statusCode == SUCCESS_STATUS || statusCode == SUCCESS_WITH_INFO_STATUS) {
+      if (statusCode == TStatusCode.SUCCESS_STATUS
+          || statusCode == TStatusCode.SUCCESS_WITH_INFO_STATUS) {
         resultSet =
             getResultSetResp(response.getStatus(), operationHandle, response.toString(), -1, true);
       }
@@ -512,4 +512,48 @@ private TCLIService.Client createThriftClient(
 
     return new TCLIService.Client(protocol);
   }
+
+  private void checkResponseForErrors(TExecuteStatementResp response) throws SQLException {
+    if (!response.isSetOperationHandle()) {
+      throw new DatabricksSQLException("Operation handle not set");
+    }
+    if (isErrorStatusCode(response.status.statusCode)) {
+      throw new DatabricksSQLException(
+          response.status.getErrorMessage(), response.status.getSqlState());
+    }
+  }
+
+  private void checkOperationStatusForErrors(TGetOperationStatusResp statusResp)
+      throws SQLException {
+    if (statusResp != null
+        && statusResp.isSetOperationState()
+        && isErrorOperationState(statusResp.getOperationState())) {
+      throw new DatabricksSQLException("Operation state erroneous");
+    }
+  }
+
+  private boolean shouldContinuePolling(TGetOperationStatusResp statusResp) {
+    return statusResp == null
+        || !statusResp.isSetOperationState()
+        || isPendingOperationState(statusResp.getOperationState());
+  }
+
+  private boolean hasResultDataInDirectResults(TExecuteStatementResp response) {
+    return response.isSetDirectResults()
+        && response.getDirectResults().isSetResultSet()
+        && response.getDirectResults().isSetResultSetMetadata();
+  }
+
+  private boolean isErrorStatusCode(TStatusCode statusCode) {
+    return statusCode == TStatusCode.ERROR_STATUS
+        || statusCode == TStatusCode.INVALID_HANDLE_STATUS;
+  }
+
+  private boolean isErrorOperationState(TOperationState state) {
+    return state == TOperationState.ERROR_STATE || state == TOperationState.CLOSED_STATE;
+  }
+
+  private boolean isPendingOperationState(TOperationState state) {
+    return state == TOperationState.RUNNING_STATE || state == TOperationState.PENDING_STATE;
+  }
 }

src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java

@@ -10,6 +10,7 @@
 import com.databricks.jdbc.api.internal.IDatabricksStatementInternal;
 import com.databricks.jdbc.common.IDatabricksComputeResource;
 import com.databricks.jdbc.common.StatementType;
+import com.databricks.jdbc.common.util.DriverUtil;
 import com.databricks.jdbc.dbclient.IDatabricksClient;
 import com.databricks.jdbc.dbclient.IDatabricksMetadataClient;
 import com.databricks.jdbc.dbclient.impl.common.MetadataResultSetBuilder;
@@ -140,6 +141,10 @@ public DatabricksResultSet executeStatement(
             .setQueryTimeout(parentStatement.getStatement().getQueryTimeout())
             .setCanReadArrowResult(this.connectionContext.shouldEnableArrow())
             .setCanDownloadResult(true);
+    if (!DriverUtil.isRunningAgainstFake()) {
+      // run queries in async mode if not using fake services
+      request.setRunAsync(true);
+    }
     return thriftAccessor.execute(request, parentStatement, session, statementType);
   }
 

src/test/java/com/databricks/jdbc/TestConstants.java

@@ -43,7 +43,6 @@ public class TestConstants {
       "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=http;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UserAgentEntry=TEST/24.2.0.2712019";
   public static final String CLUSTER_JDBC_URL =
       "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=http;ssl=1;httpPath=sql/protocolv1/o/6051921418418893/1115-130834-ms4m0yv;AuthMech=3;UserAgentEntry=MyApp";
-
   public static final TRowSet binaryRowSet =
       new TRowSet()
           .setColumns(
@@ -83,7 +82,6 @@ public class TestConstants {
           .setColumns(
               Collections.singletonList(
                   TColumn.i64Val(new TI64Column().setValues(List.of(1L, 5L)))));
-
   public static final TRowSet stringRowSet =
       new TRowSet()
           .setColumns(
@@ -163,7 +161,6 @@ public class TestConstants {
       "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;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=2;CFProxyUID=cfProxyUser;CFProxyPwd=cfProxyPassword;";
   public 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";
-
   public static final List<TSparkArrowBatch> ARROW_BATCH_LIST =
       Collections.singletonList(
           new TSparkArrowBatch().setRowCount(0).setBatch(new byte[] {65, 66, 67}));