SNOW-980046 Implement queryStatusV2

src/main/java/net/snowflake/client/core/SFBaseSession.java

@@ -819,10 +819,19 @@ public Object getSessionPropertyByKey(String propertyName) {
   /**
    * @param queryID query ID of the query whose status is being investigated
    * @return enum of type QueryStatus indicating the query's status
+   * @deprecated Use {@link #getQueryStatusV2(String)}
    * @throws SQLException
    */
+  @Deprecated
   public abstract QueryStatus getQueryStatus(String queryID) throws SQLException;
 
+  /**
+   * @param queryID query ID of the query whose status is being investigated
+   * @return QueryStatusV2 indicating the query's status
+   * @throws SQLException
+   */
+  public abstract QueryStatusV2 getQueryStatusV2(String queryID) throws SQLException;
+
   /**
    * Validates the connection properties used by this session, and returns a list of missing
    * properties.

src/main/java/net/snowflake/client/core/SFSession.java

@@ -164,13 +164,7 @@ public void addQueryToActiveQueryList(String queryID) {
     activeAsyncQueries.add(queryID);
   }
 
-  /**
-   * @param queryID query ID of the query whose status is being investigated
-   * @return enum of type QueryStatus indicating the query's status
-   * @throws SQLException
-   */
-  @Override
-  public QueryStatus getQueryStatus(String queryID) throws SQLException {
+  private JsonNode getQueryMetadata(String queryID) throws SQLException {
     // create the URL to check the query monitoring endpoint
     String statusUrl = "";
     String sessionUrl = getUrl();
@@ -245,7 +239,18 @@ else if (ex instanceof SFException) {
         }
       }
     } while (sessionRenewed);
-    JsonNode queryNode = jsonNode.path("data").path("queries");
+    return jsonNode.path("data").path("queries");
+  }
+
+  /**
+   * @param queryID query ID of the query whose status is being investigated
+   * @return enum of type QueryStatus indicating the query's status
+   * @throws SQLException
+   * @deprecated the returned enum is error-prone, use {@link #getQueryStatusV2} instead
+   */
+  @Deprecated
+  public QueryStatus getQueryStatus(String queryID) throws SQLException {
+    JsonNode queryNode = getQueryMetadata(queryID);
     String queryStatus = "";
     String errorMessage = "";
     int errorCode = 0;
@@ -284,6 +289,54 @@ else if (isAnError(result)) {
     return result;
   }
 
+  /**
+   * @param queryID query ID of the query whose status is being investigated
+   * @return a QueryStatusV2 instance indicating the query's status
+   * @throws SQLException
+   */
+  public QueryStatusV2 getQueryStatusV2(String queryID) throws SQLException {
+    JsonNode queryNode = getQueryMetadata(queryID);
+    logger.debug("Query status: {}", queryNode.asText());
+    if (queryNode.isEmpty()) {
+      return QueryStatusV2.empty();
+    }
+    JsonNode node = queryNode.get(0);
+    long endTime = node.path("endTime").asLong(0);
+    int errorCode = node.path("errorCode").asInt(0);
+    String errorMessage = node.path("errorMessage").asText("No error reported");
+    String id = node.path("id").asText("");
+    String name = node.path("status").asText("");
+    long sessionId = node.path("sessionId").asLong(0);
+    String sqlText = node.path("sqlText").asText("");
+    long startTime = node.path("startTime").asLong(0);
+    String state = node.path("state").asText("");
+    int totalDuration = node.path("totalDuration").asInt(0);
+    String warehouseExternalSize = node.path("warehouseExternalSize").asText(null);
+    int warehouseId = node.path("warehouseId").asInt(0);
+    String warehouseName = node.path("warehouseName").asText(null);
+    String warehouseServerType = node.path("warehouseServerType").asText(null);
+    QueryStatusV2 result =
+        new QueryStatusV2(
+            endTime,
+            errorCode,
+            errorMessage,
+            id,
+            name,
+            sessionId,
+            sqlText,
+            startTime,
+            state,
+            totalDuration,
+            warehouseExternalSize,
+            warehouseId,
+            warehouseName,
+            warehouseServerType);
+    if (!result.isStillRunning()) {
+      activeAsyncQueries.remove(queryID);
+    }
+    return result;
+  }
+
   /**
    * Add a property If a property is known for connection, add it to connection properties If not,
    * add it as a dynamic session parameters

src/main/java/net/snowflake/client/core/SFStatement.java

@@ -10,6 +10,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.*;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -649,8 +650,8 @@ public long getConservativeMemoryLimit() {
    */
   @Override
   public String[] getChildQueryIds(String queryID) throws SQLException {
-    QueryStatus qs = session.getQueryStatus(queryID);
-    if (QueryStatus.isStillRunning(qs)) {
+    QueryStatusV2 qs = session.getQueryStatusV2(queryID);
+    if (qs.isStillRunning()) {
       throw new SQLException(
           "Status of query associated with resultSet is "
               + qs.getDescription()

src/main/java/net/snowflake/client/jdbc/QueryStatusV2.java

@@ -0,0 +1,138 @@
+package net.snowflake.client.jdbc;
+
+import net.snowflake.client.core.QueryStatus;
+
+public final class QueryStatusV2 {
+  private final long endTime;
+  private final int errorCode;
+  private final String errorMessage;
+  private final String id;
+  private final String name;
+  private final long sessionId;
+  private final String sqlText;
+  private final long startTime;
+  private final String state;
+  private final QueryStatus status;
+  private final int totalDuration;
+  private final String warehouseExternalSize;
+  private final int warehouseId;
+  private final String warehouseName;
+  private final String warehouseServerType;
+
+  public QueryStatusV2(
+      long endTime,
+      int errorCode,
+      String errorMessage,
+      String id,
+      String name,
+      long sessionId,
+      String sqlText,
+      long startTime,
+      String state,
+      int totalDuration,
+      String warehouseExternalSize,
+      int warehouseId,
+      String warehouseName,
+      String warehouseServerType) {
+    this.endTime = endTime;
+    this.errorCode = errorCode;
+    this.errorMessage = errorMessage;
+    this.id = id;
+    this.name = name;
+    this.sessionId = sessionId;
+    this.sqlText = sqlText;
+    this.startTime = startTime;
+    this.state = state;
+    this.status = QueryStatus.getStatusFromString(name);
+    this.totalDuration = totalDuration;
+    this.warehouseExternalSize = warehouseExternalSize;
+    this.warehouseId = warehouseId;
+    this.warehouseName = warehouseName;
+    this.warehouseServerType = warehouseServerType;
+  }
+
+  public static QueryStatusV2 empty() {
+    return new QueryStatusV2(0, 0, "", "", "", 0, "", 0, "", 0, "", 0, "", "");
+  }
+
+  public boolean isEmpty() {
+    return name.isEmpty();
+  }
+
+  public boolean isStillRunning() {
+    return QueryStatus.isStillRunning(status);
+  }
+
+  public boolean isSuccess() {
+    return status == QueryStatus.SUCCESS;
+  }
+
+  public boolean isAnError() {
+    return QueryStatus.isAnError(status);
+  }
+
+  public long getEndTime() {
+    return endTime;
+  }
+
+  public int getErrorCode() {
+    return errorCode;
+  }
+
+  public String getErrorMessage() {
+    return errorMessage;
+  }
+
+  public String getId() {
+    return id;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public long getSessionId() {
+    return sessionId;
+  }
+
+  public String getSqlText() {
+    return sqlText;
+  }
+
+  public long getStartTime() {
+    return startTime;
+  }
+
+  public String getState() {
+    return state;
+  }
+
+  public int getTotalDuration() {
+    return totalDuration;
+  }
+
+  public String getWarehouseExternalSize() {
+    return warehouseExternalSize;
+  }
+
+  public int getWarehouseId() {
+    return warehouseId;
+  }
+
+  public String getWarehouseName() {
+    return warehouseName;
+  }
+
+  public String getWarehouseServerType() {
+    return warehouseServerType;
+  }
+
+  /** To preserve compatibility with {@link QueryStatus} */
+  public String getDescription() {
+    return name;
+  }
+
+  public QueryStatus getStatus() {
+    return status;
+  }
+}

src/main/java/net/snowflake/client/jdbc/SFAsyncResultSet.java

@@ -28,6 +28,7 @@ public class SFAsyncResultSet extends SnowflakeBaseResultSet
   private SFBaseSession session;
   private Statement extraStatement;
   private QueryStatus lastQueriedStatus = NO_DATA;
+  private QueryStatusV2 lastQueriedStatusV2 = QueryStatusV2.empty();
 
   /**
    * Constructor takes an inputstream from the API response that we get from executing a SQL
@@ -100,9 +101,9 @@ public QueryStatus getStatus() throws SQLException {
     if (this.lastQueriedStatus == QueryStatus.SUCCESS) {
       return this.lastQueriedStatus;
     }
-    this.lastQueriedStatus = session.getQueryStatus(this.queryID);
     // if query has completed successfully, cache its success status to avoid unnecessary future
     // server calls
+    this.lastQueriedStatus = session.getQueryStatus(this.queryID);
     return this.lastQueriedStatus;
   }
 
@@ -111,6 +112,23 @@ public String getQueryErrorMessage() throws SQLException {
     return this.lastQueriedStatus.getErrorMessage();
   }
 
+  @Override
+  public QueryStatusV2 getStatusV2() throws SQLException {
+    if (session == null) {
+      throw new SQLException("Session not set");
+    }
+    if (this.queryID == null) {
+      throw new SQLException("QueryID unknown");
+    }
+    if (this.lastQueriedStatusV2.isSuccess()) {
+      return this.lastQueriedStatusV2;
+    }
+    this.lastQueriedStatusV2 = session.getQueryStatusV2(this.queryID);
+    // if query has completed successfully, cache its metadata to avoid unnecessary future server
+    // calls
+    return this.lastQueriedStatusV2;
+  }
+
   /**
    * helper function for next() and getMetaData(). Calls result_scan to get resultSet after
    * asynchronous query call

src/main/java/net/snowflake/client/jdbc/SnowflakeResultSet.java

@@ -37,6 +37,18 @@ public interface SnowflakeResultSet {
    */
   String getQueryErrorMessage() throws SQLException;
 
+  /**
+   * This function retrieves the status of an asynchronous query. An empty ResultSet object has
+   * already been returned, but the query may still be running. This function can be used to query
+   * whether it is possible to retrieve results from the ResultSet already.
+   *
+   * <p><code>status.isSuccess()</code> means that results can be retrieved.
+   *
+   * @return an instance containing query metadata
+   * @throws SQLException
+   */
+  QueryStatusV2 getStatusV2() throws SQLException;
+
   /**
    * Get a list of ResultSetSerializables for the ResultSet in order to parallel processing
    *

src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetV1.java

@@ -50,6 +50,18 @@ public QueryStatus getStatus() throws SQLException {
     throw new SnowflakeLoggedFeatureNotSupportedException(session);
   }
 
+  /**
+   * This function is not supported for synchronous queries
+   *
+   * @return no return value; exception is always thrown
+   * @throws SQLFeatureNotSupportedException
+   */
+  @Override
+  public QueryStatusV2 getStatusV2() throws SQLException {
+    throw new SnowflakeLoggedFeatureNotSupportedException(
+        session, "This function is only supported for asynchronous queries.");
+  }
+
   /**
    * This function is not supported for synchronous queries
    *