fix: Timeout custom exception (#108)

firebolt-db · Sep 25, 2024 · db54a7a · db54a7a
1 parent 785cc62
commit db54a7a
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 12 deletions.
diff --git a/FireboltDotNetSdk.Tests/Integration/ExecuteTest.cs b/FireboltDotNetSdk.Tests/Integration/ExecuteTest.cs
@@ -126,11 +126,12 @@ public void WithTimeout(int? timeout, bool async)
 
             if (async)
             {
-                Assert.ThrowsAsync<AggregateException>(command.ExecuteReaderAsync);
+                var aggEx = Assert.ThrowsAsync<AggregateException>(command.ExecuteReaderAsync);
+                Assert.That(aggEx?.InnerException, Is.InstanceOf<FireboltTimeoutException>());
             }
             else
             {
-                Assert.Throws<TaskCanceledException>(() => command.ExecuteReader());
+                Assert.Throws<FireboltTimeoutException>(() => command.ExecuteReader());
             }
         }
 
@@ -153,6 +154,21 @@ public async Task WithZeroTimeout(bool async)
             }
         }
 
+        [Test]
+        public void WithTokenCancel()
+        {
+            using var conn = new FireboltConnection(SYSTEM_CONNECTION_STRING);
+            conn.Open();
+            DbCommand command = conn.CreateCommand();
+            command.CommandText = LONG_QUERY;
+            CancellationTokenSource source = new CancellationTokenSource();
+            CancellationToken token = source.Token;
+            var task = command.ExecuteReaderAsync(token);
+            source.Cancel();
+            var aggEx = Assert.ThrowsAsync<AggregateException>(async () => await task);
+            Assert.That(aggEx?.InnerException, Is.InstanceOf<TaskCanceledException>());
+        }
+
         [Test]
         [Category("general")]
         public void ShortTimeout()

diff --git a/FireboltNETSDK/Client/FireboltCommand.cs b/FireboltNETSDK/Client/FireboltCommand.cs
@@ -45,6 +45,7 @@ public class FireboltCommand : DbCommand
         private static readonly string FORBIDDEN_PROPERTY_ERROR_SET_SUFFIX = "Try again with a different parameter name.";
         private static readonly string USE_ERROR = FORBIDDEN_PROPERTY_ERROR_PREFIX + FORBIDDEN_PROPERTY_ERROR_USE_SUFFIX;
         private static readonly string SET_ERROR = FORBIDDEN_PROPERTY_ERROR_PREFIX + FORBIDDEN_PROPERTY_ERROR_SET_SUFFIX;
+        private static readonly TimeSpan regexTimeout = TimeSpan.FromSeconds(5);
         internal static readonly string BYTE_ARRAY_PREFIX = "\\x";
 
         private FireboltConnection? _connection;
@@ -149,8 +150,7 @@ public override int CommandTimeout
 
         internal QueryResult? Execute(string commandText)
         {
-            CancellationTokenSource cancellationTokenSource = CommandTimeoutMillis == 0 ? new CancellationTokenSource() : new CancellationTokenSource(CommandTimeoutMillis);
-            return CreateQueryResult(ExecuteCommandAsync(commandText, cancellationTokenSource.Token).GetAwaiter().GetResult());
+            return CreateQueryResult(ExecuteCommandAsyncWithCommandTimeout(commandText, CancellationToken.None).GetAwaiter().GetResult());
         }
 
         private QueryResult? CreateQueryResult(string? response)
@@ -163,6 +163,27 @@ private DbDataReader CreateDbDataReader(QueryResult? queryResult)
             return queryResult != null ? new FireboltDataReader(null, queryResult, 0) : throw new InvalidOperationException("No result produced");
         }
 
+        private async Task<string?> ExecuteCommandAsyncWithCommandTimeout(string commandText, CancellationToken cancellationToken)
+        {
+            if (CommandTimeoutMillis == 0)
+            {
+                return await ExecuteCommandAsync(commandText, cancellationToken);
+            }
+            using (var timeoutSource = new CancellationTokenSource(TimeSpan.FromMilliseconds(CommandTimeoutMillis)))
+            using (var linkedTokenSource =
+                   CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token))
+            {
+                try
+                {
+                    return await ExecuteCommandAsync(commandText, linkedTokenSource.Token);
+                }
+                catch (TaskCanceledException) when (timeoutSource.Token.IsCancellationRequested)
+                {
+                    throw new FireboltTimeoutException(CommandTimeoutMillis);
+                }
+            }
+        }
+
         private async Task<string?> ExecuteCommandAsync(string commandText, CancellationToken cancellationToken)
         {
             if (Connection == null)
@@ -219,13 +240,8 @@ protected override Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBehavior b
             {
                 cancellationToken.Register(Cancel);
             }
-            var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-            if (CommandTimeoutMillis != 0)
-            {
-                // Add token timeout
-                linkedTokenSource.CancelAfter(CommandTimeoutMillis);
-            }
-            return ExecuteCommandAsync(StrictCommandText, linkedTokenSource.Token).ContinueWith(result => CreateDbDataReader(CreateQueryResult(result.Result)));
+            return ExecuteCommandAsyncWithCommandTimeout(StrictCommandText, cancellationToken).ContinueWith(
+                result => CreateDbDataReader(CreateQueryResult(result.Result)));
         }
 
         public override Task<object?> ExecuteScalarAsync(CancellationToken cancellationToken)
@@ -245,7 +261,7 @@ private string GetParamQuery(string commandText)
                 {
                     string pattern = string.Format(@"\{0}\b", parameter.ParameterName);
                     string verifyParameters = GetParamValue(parameter.Value);
-                    commandText = Regex.Replace(commandText, pattern, verifyParameters, RegexOptions.IgnoreCase);
+                    commandText = Regex.Replace(commandText, pattern, verifyParameters, RegexOptions.IgnoreCase, regexTimeout);
                 }
                 return commandText;
             }

diff --git a/FireboltNETSDK/Exception/FireboltTimeoutException.cs b/FireboltNETSDK/Exception/FireboltTimeoutException.cs
@@ -0,0 +1,8 @@
+namespace FireboltDotNetSdk.Exception;
+
+public class FireboltTimeoutException : FireboltException
+{
+    public FireboltTimeoutException(int timeoutMillis) : base($"Query execution timeout. The query did not complete within {timeoutMillis} milliseconds.")
+    {
+    }
+}