Add util functions for user to check query status is still running or…

… an error
snowflakedb · Mar 25, 2024 · e1c50b6 · e1c50b6
1 parent 73358b1
commit e1c50b6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
@@ -197,7 +197,8 @@ public async Task TestAsyncExecQueryAsync()
                     var queryStatus = await cmd.GetQueryStatusAsync(queryId, CancellationToken.None).ConfigureAwait(false);
 
                     // Assert
-                    Assert.IsTrue(QueryStatusExtensions.IsStillRunning(queryStatus));
+                    Assert.IsTrue(conn.IsStillRunning(queryStatus));
+                    Assert.IsFalse(conn.IsAnError(queryStatus));
 
                     // Act
                     DbDataReader reader = await cmd.GetResultsFromQueryIdAsync(queryId, CancellationToken.None).ConfigureAwait(false);
@@ -237,7 +238,7 @@ public async Task TestExecuteNormalQueryWhileAsyncExecQueryIsRunningAsync()
                 var queryStatus = await cmd.GetQueryStatusAsync(queryId, CancellationToken.None).ConfigureAwait(false);
 
                 // Assert
-                Assert.IsTrue(QueryStatusExtensions.IsStillRunning(queryStatus));
+                Assert.IsTrue(connections[0].IsStillRunning(queryStatus));
             }
 
             // Execute a normal query
@@ -291,7 +292,7 @@ public async Task TestAsyncExecCancelWhileGettingResultsAsync()
                     var queryStatus = await cmd.GetQueryStatusAsync(queryId, CancellationToken.None).ConfigureAwait(false);
 
                     // Assert
-                    Assert.IsTrue(QueryStatusExtensions.IsStillRunning(queryStatus));
+                    Assert.IsTrue(conn.IsStillRunning(queryStatus));
 
                     // Act
                     cancelToken.Cancel();
@@ -324,7 +325,7 @@ public async Task TestFailedAsyncExecQueryThrowsErrorAsync()
                     // Act
                     var queryId = await cmd.ExecuteAsyncInAsyncMode(CancellationToken.None).ConfigureAwait(false);
                     var queryStatus = await cmd.GetQueryStatusAsync(queryId, CancellationToken.None).ConfigureAwait(false);
-                    while (statusRetryCount < statusMaxRetryCount && QueryStatusExtensions.IsStillRunning(queryStatus))
+                    while (statusRetryCount < statusMaxRetryCount && conn.IsStillRunning(queryStatus))
                     {
                         Thread.Sleep(1000);
                         queryStatus = await cmd.GetQueryStatusAsync(queryId, CancellationToken.None).ConfigureAwait(false);
@@ -1351,9 +1352,11 @@ public void TestAsyncExecQuery()
 
                     // Act
                     queryId = cmd.ExecuteInAsyncMode();
+                    var queryStatus = cmd.GetQueryStatus(queryId);
 
                     // Assert
-                    Assert.IsTrue(QueryStatusExtensions.IsStillRunning(cmd.GetQueryStatus(queryId)));
+                    Assert.IsTrue(conn.IsStillRunning(queryStatus));
+                    Assert.IsFalse(conn.IsAnError(queryStatus));
 
                     // Act
                     DbDataReader reader = cmd.GetResultsFromQueryId(queryId);
@@ -1391,7 +1394,7 @@ public void TestExecuteNormalQueryWhileAsyncExecQueryIsRunning()
                 queryId = cmd.ExecuteInAsyncMode();
 
                 // Assert
-                Assert.IsTrue(QueryStatusExtensions.IsStillRunning(cmd.GetQueryStatus(queryId)));
+                Assert.IsTrue(connections[0].IsStillRunning(cmd.GetQueryStatus(queryId)));
             }
 
             // Execute a normal query
@@ -1444,7 +1447,7 @@ public void TestFailedAsyncExecQueryThrowsError()
 
                     // Act
                     queryId = cmd.ExecuteInAsyncMode();
-                    while (statusRetryCount < statusMaxRetryCount && QueryStatusExtensions.IsStillRunning(cmd.GetQueryStatus(queryId)))
+                    while (statusRetryCount < statusMaxRetryCount && conn.IsStillRunning(cmd.GetQueryStatus(queryId)))
                     {
                         Thread.Sleep(1000);
                         statusRetryCount++;

diff --git a/Snowflake.Data/Client/SnowflakeDbConnection.cs b/Snowflake.Data/Client/SnowflakeDbConnection.cs
@@ -402,6 +402,16 @@ internal void registerConnectionCancellationCallback(CancellationToken externalC
             }
         }
 
+        public bool IsStillRunning(QueryStatus status)
+        {
+            return QueryStatusExtensions.IsStillRunning(status);
+        }
+
+        public bool IsAnError(QueryStatus status)
+        {
+            return QueryStatusExtensions.IsAnError(status);
+        }
+
         ~SnowflakeDbConnection()
         {
             Dispose(false);