SNOW-668836: Fix ConnectionLatestIt.testQueryStatusErrorMessageAndErrorCodeChangeOnAsyncQuery (#1644)

Author: sfc-gh-dprzybysz

TestOnly/pom.xml

@@ -13,6 +13,7 @@
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
     <arrow.version>10.0.1</arrow.version>
+    <awaitility.version>4.2.0</awaitility.version>
     <jacksondatabind.version>2.13.4.2</jacksondatabind.version>
     <jacoco.version>0.8.4</jacoco.version>
     <jacoco.skip.instrument>true</jacoco.skip.instrument>
@@ -59,6 +60,12 @@
       <version>${mockito.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <version>${awaitility.version}</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>commons-dbcp</groupId>
       <artifactId>commons-dbcp</artifactId>

parent-pom.xml

@@ -22,6 +22,7 @@
     <arrow.version>10.0.1</arrow.version>
     <asm.version>9.3</asm.version>
     <avro.version>1.8.1</avro.version>
+    <awaitility.version>4.2.0</awaitility.version>
     <awssdk.version>1.12.501</awssdk.version>
     <azure.storage.version>5.0.0</azure.storage.version>
     <bouncycastle.version>1.74</bouncycastle.version>
@@ -475,6 +476,12 @@
         <version>${mockito.version}</version>
         <scope>test</scope>
       </dependency>
+      <dependency>
+        <groupId>org.awaitility</groupId>
+        <artifactId>awaitility</artifactId>
+        <version>${awaitility.version}</version>
+        <scope>test</scope>
+      </dependency>
     </dependencies>
   </dependencyManagement>
 
@@ -704,5 +711,9 @@
       <groupId>org.mockito</groupId>
       <artifactId>mockito-inline</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+    </dependency>
   </dependencies>
 </project>

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

@@ -37,18 +37,34 @@ public String getDescription() {
     return this.description;
   }
 
+  /**
+   * @deprecated use {@link net.snowflake.client.jdbc.QueryStatusV2} instead
+   */
+  @Deprecated
   public String getErrorMessage() {
     return this.errorMessage;
   }
 
+  /**
+   * @deprecated use {@link net.snowflake.client.jdbc.QueryStatusV2} instead
+   */
+  @Deprecated
   public int getErrorCode() {
     return this.errorCode;
   }
 
+  /**
+   * @deprecated use {@link net.snowflake.client.jdbc.QueryStatusV2} instead
+   */
+  @Deprecated
   public void setErrorMessage(String message) {
     this.errorMessage = message;
   }
 
+  /**
+   * @deprecated use {@link net.snowflake.client.jdbc.QueryStatusV2} instead
+   */
+  @Deprecated
   public void setErrorCode(int errorCode) {
     this.errorCode = errorCode;
   }

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

@@ -7,6 +7,7 @@
 import static net.snowflake.client.jdbc.ConnectionIT.BAD_REQUEST_GS_CODE;
 import static net.snowflake.client.jdbc.ConnectionIT.INVALID_CONNECTION_INFO_CODE;
 import static net.snowflake.client.jdbc.ConnectionIT.WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS;
+import static org.awaitility.Awaitility.await;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -390,36 +391,24 @@ public void testAsyncAndSynchronousQueries() throws SQLException {
     con.close();
   }
 
-  /**
-   * Tests that error message and error code are reset after an error. This test is not reliable as
-   * it uses sleep() call. It works locally but failed with PR.
-   *
-   * @throws SQLException
-   * @throws InterruptedException
-   */
-  // @Test
-  public void testQueryStatusErrorMessageAndErrorCode() throws SQLException, InterruptedException {
-    // open connection and run asynchronous query
-    Connection con = getConnection();
-    Statement statement = con.createStatement();
-    statement.execute("create or replace table testTable(colA string, colB boolean)");
-    statement.execute("insert into testTable values ('test', true)");
-    ResultSet rs1 =
-        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select * from testTable");
-    QueryStatus status = rs1.unwrap(SnowflakeResultSet.class).getStatus();
-    // Set the error message and error code so we can confirm they are reset when getStatus() is
-    // called.
-    status.setErrorMessage(QueryStatus.FAILED_WITH_ERROR.toString());
-    status.setErrorCode(2003);
-    Thread.sleep(300);
-    status = rs1.unwrap(SnowflakeResultSet.class).getStatus();
-    // Assert status of query is a success
-    assertEquals(QueryStatus.SUCCESS, status);
-    assertEquals("No error reported", status.getErrorMessage());
-    assertEquals(0, status.getErrorCode());
-    statement.execute("drop table if exists testTable");
-    statement.close();
-    con.close();
+  /** Can be used in > 3.14.4 (when {@link QueryStatusV2} was added) */
+  @Test
+  public void testQueryStatusErrorMessageAndErrorCodeChangeOnAsyncQuery() throws SQLException {
+    try (Connection con = getConnection();
+        Statement statement = con.createStatement();
+        ResultSet rs1 =
+            statement
+                .unwrap(SnowflakeStatement.class)
+                .executeAsyncQuery("select count(*) from table(generator(timeLimit => 2))")) {
+      SnowflakeResultSet sfResultSet = rs1.unwrap(SnowflakeResultSet.class);
+      // status should change state to RUNNING and then to SUCCESS
+      await()
+          .atMost(Duration.ofSeconds(5))
+          .until(() -> sfResultSet.getStatusV2().getStatus(), equalTo(QueryStatus.RUNNING));
+      await()
+          .atMost(Duration.ofSeconds(5))
+          .until(() -> sfResultSet.getStatusV2().getStatus(), equalTo(QueryStatus.SUCCESS));
+    }
   }
 
   @Test