SNOW-538588 refactor to support async in stored proc (#1324)

* SNOW-538588 refactor to support async in stored proc

* remove unwanted change

---------

Co-authored-by: Zihan Li <zihan.li@snowflake.com>

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

@@ -9,6 +9,7 @@
 
 import com.google.common.base.Strings;
 import java.sql.DriverPropertyInfo;
+import java.sql.SQLException;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -793,6 +794,13 @@ public Object getSessionPropertyByKey(String propertyName) {
    */
   public abstract boolean isSafeToClose();
 
+  /**
+   * @param queryID query ID of the query whose status is being investigated
+   * @return enum of type QueryStatus indicating the query's status
+   * @throws SQLException
+   */
+  public abstract QueryStatus getQueryStatus(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

@@ -162,6 +162,7 @@ public void addQueryToActiveQueryList(String queryID) {
    * @return enum of type QueryStatus indicating the query's status
    * @throws SQLException
    */
+  @Override
   public QueryStatus getQueryStatus(String queryID) throws SQLException {
     // create the URL to check the query monitoring endpoint
     String statusUrl = "";

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

@@ -14,6 +14,7 @@
 import java.util.regex.Pattern;
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFBaseResultSet;
+import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SFSession;
 import net.snowflake.common.core.SqlState;
 
@@ -23,7 +24,7 @@ class SFAsyncResultSet extends SnowflakeBaseResultSet implements SnowflakeResult
   private ResultSet resultSetForNext = new SnowflakeResultSetV1.EmptyResultSet();
   private boolean resultSetForNextInitialized = false;
   private String queryID;
-  private SFSession session;
+  private SFBaseSession session;
   private Statement extraStatement;
   private QueryStatus lastQueriedStatus = NO_DATA;
 
@@ -42,7 +43,7 @@ class SFAsyncResultSet extends SnowflakeBaseResultSet implements SnowflakeResult
     super(statement);
     this.sfBaseResultSet = sfBaseResultSet;
     this.queryID = sfBaseResultSet.getQueryId();
-    this.session = (SFSession) sfBaseResultSet.getSession();
+    this.session = sfBaseResultSet.getSession();
     this.extraStatement = statement;
     this.resultSetMetaData = new SnowflakeResultSetMetaDataV1(sfBaseResultSet.getMetaData());
     this.resultSetMetaData.setQueryIdForAsyncResults(this.queryID);

src/test/java/net/snowflake/client/jdbc/MockConnectionTest.java

@@ -670,6 +670,11 @@ public List<DriverPropertyInfo> checkProperties() {
     @Override
     public void close() {}
 
+    @Override
+    public QueryStatus getQueryStatus(String queryID) {
+      return null;
+    }
+
     @Override
     public Telemetry getTelemetryClient() {
       return new Telemetry() {