Add support for command cancellation

Giorgi · Jul 21, 2024 · fe02f4c · fe02f4c
1 parent 2c78e20
commit fe02f4c
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 4 deletions.
diff --git a/DuckDB.NET.Bindings/NativeMethods/NativeMethods.Startup.cs b/DuckDB.NET.Bindings/NativeMethods/NativeMethods.Startup.cs
@@ -29,6 +29,9 @@ public static class Startup
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_disconnect")]
         public static extern void DuckDBDisconnect(out IntPtr connection);
 
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_interrupt")]
+        public static extern void DuckDBInterrupt(IntPtr connection);
+
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_library_version")]
         public static extern IntPtr DuckDBLibraryVersion();
     }

diff --git a/DuckDB.NET.Data/DuckDBCommand.cs b/DuckDB.NET.Data/DuckDBCommand.cs
@@ -70,7 +70,12 @@ public DuckDBCommand(string commandText, DuckDBConnection connection)
     }
 
     public override void Cancel()
-    { }
+    {
+        if (connection != null)
+        {
+            NativeMethods.Startup.DuckDBInterrupt(connection.NativeConnection.DangerousGetHandle());
+        }
+    }
 
     public override int ExecuteNonQuery()
     {

diff --git a/DuckDB.NET.Data/Internal/PreparedStatement.cs b/DuckDB.NET.Data/Internal/PreparedStatement.cs
@@ -87,7 +87,18 @@ public DuckDBResult Execute(DuckDBParameterCollection parameterCollection, bool
         {
             var errorMessage = NativeMethods.Query.DuckDBResultError(ref queryResult).ToManagedString(false);
             queryResult.Dispose();
-            throw new DuckDBException(string.IsNullOrEmpty(errorMessage) ? "DuckDBQuery failed" : errorMessage, status);
+
+            if (string.IsNullOrEmpty(errorMessage))
+            {
+                errorMessage = "DuckDB execution failed";
+            }
+
+            if (errorMessage.StartsWith("INTERRUPT Error"))
+            {
+                throw new OperationCanceledException();
+            }
+
+            throw new DuckDBException(errorMessage, status);
         }
 
         return queryResult;

diff --git a/DuckDB.NET.Test/DuckDBDataReaderTests.cs b/DuckDB.NET.Test/DuckDBDataReaderTests.cs
@@ -5,6 +5,8 @@
 using System.Data;
 using DuckDB.NET.Native;
 using Xunit;
+using System.Threading.Tasks;
+using System.Threading;
 
 namespace DuckDB.NET.Test;
 
@@ -356,4 +358,20 @@ public void ReadDecimalSchemaWithoutTableRow()
         schemaTable.Rows[0]["NumericScale"].Should().Be(0);
         schemaTable.Rows[0]["NumericPrecision"].Should().Be(0);
     }
+
+    [Fact]
+    public async Task CancellingLongRunningQueryThrowsOperationCancelledException()
+    {
+        Command.CommandText = @"create table cnt as WITH RECURSIVE
+                       cnt(x) AS (
+                          SELECT 1
+                          UNION ALL
+                          SELECT x+1 FROM cnt
+                           where x < 300000
+                    ) select * from cnt;";
+
+        var source = new CancellationTokenSource(1000);
+
+        await Command.Invoking(async c => await c.ExecuteReaderAsync(source.Token)).Should().ThrowAsync<OperationCanceledException>();
+    }
 }
diff --git a/DuckDB.NET.Test/ExecuteNonQueryTests.cs b/DuckDB.NET.Test/ExecuteNonQueryTests.cs
@@ -1,5 +1,4 @@
-using DuckDB.NET.Data;
-using FluentAssertions;
+using FluentAssertions;
 using System.Collections.Generic;
 using Xunit;