Re-add check if token is cancelled

snowflakedb · Mar 21, 2024 · 7185f79 · 7185f79
1 parent 21915dc
commit 7185f79
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs b/Snowflake.Data.Tests/IntegrationTests/SFDbCommandIT.cs
@@ -295,20 +295,11 @@ public async Task TestAsyncExecCancelWhileGettingResultsAsync()
 
                     // Act
                     cancelToken.Cancel();
-#if NET471
                     var thrown = Assert.ThrowsAsync<OperationCanceledException>(async () =>
                         await cmd.GetResultsFromQueryIdAsync(queryId, cancelToken.Token).ConfigureAwait(false));
-#else
-                    var thrown = Assert.ThrowsAsync<TaskCanceledException>(async () =>
-                        await cmd.GetResultsFromQueryIdAsync(queryId, cancelToken.Token).ConfigureAwait(false));
-#endif
 
                     // Assert
-#if NET472
-                    Assert.IsTrue(thrown.Message.Contains("A task was canceled"));
-#else
                     Assert.IsTrue(thrown.Message.Contains("The operation was canceled"));
-#endif
                 }
 
                 await conn.CloseAsync(CancellationToken.None).ConfigureAwait(false);

diff --git a/Snowflake.Data/Client/SnowflakeDbCommand.cs b/Snowflake.Data/Client/SnowflakeDbCommand.cs
@@ -366,6 +366,12 @@ internal async Task RetryUntilQueryResultIsAvailable(string queryId, Cancellatio
             QueryStatus status;
             while (true)
             {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    logger.Debug("Cancellation requested for getting results from query id");
+                    cancellationToken.ThrowIfCancellationRequested();
+                }
+
                 status = isAsync ? await GetQueryStatusAsync(queryId, cancellationToken) : GetQueryStatus(queryId);
 
                 if (!QueryStatuses.IsStillRunning(status))