Add Invoke-Sql cancellation support.

Author: sharpjs

PSql.Engine.Tests/Integration/SqlConnectionIntegrationTests.cs

@@ -12,15 +12,15 @@ namespace PSql.Integration;
 public class SqlConnectionIntegrationTests
 {
     [Test]
-    public void ExecuteAndProjectTo_ClrTypes()
+    public void ExecuteAndProject_ClrTypes()
     {
         using var connection = new SqlConnection(
             IntegrationTestsSetup.Database.ConnectionString,
             IntegrationTestsSetup.Credential,
             TestSqlLogger.Instance
         );
 
-        using var result = connection.ExecuteAndProjectTo(
+        using var result = connection.ExecuteAndProject(
             """
             SELECT *, 10 FROM (VALUES (N'a', 1), (N'b', 2)) AS T (S, X);
             SELECT *, 20 FROM (VALUES (N'c', 3), (N'd', 4)) AS T (S, Y);
@@ -37,7 +37,7 @@ public void ExecuteAndProjectTo_ClrTypes()
 
     [Test]
     [SetCulture("kl-GL")] // Greenlandic
-    public void ExecuteAndProjectTo_SqlTypes()
+    public void ExecuteAndProject_SqlTypes()
     {
         // NOTE: It appears that the current .NET culture is what determines
         // the collation of a SqlString.  Even a collation specified explicitly
@@ -55,7 +55,7 @@ static SqlString Greenlandic(string s)
             TestSqlLogger.Instance
         );
 
-        using var result = connection.ExecuteAndProjectTo(
+        using var result = connection.ExecuteAndProject(
             """
             SELECT *, 10 FROM (VALUES (N'a', 1), (N'b', 2)) AS T (S, X);
             SELECT *, 20 FROM (VALUES
@@ -75,15 +75,15 @@ static SqlString Greenlandic(string s)
     }
 
     [Test]
-    public void ExecuteAndProjectTo_Exception()
+    public void ExecuteAndProject_Exception()
     {
         using var connection = new SqlConnection(
             IntegrationTestsSetup.Database.ConnectionString,
             IntegrationTestsSetup.Credential,
             TestSqlLogger.Instance
         );
 
-        using var result = connection.ExecuteAndProjectTo(
+        using var result = connection.ExecuteAndProject(
             """
             SELECT * FROM (VALUES (1/1)) AS T (X);
             SELECT * FROM (VALUES (1/0)) AS T (X);

PSql.Engine.Tests/SqlConnectionTests.cs

@@ -9,15 +9,15 @@ public class SqlConnectionTests
     // This test class only backfills coverage gaps in other tests.
 
     [Test]
-    public void Foo()
+    public void ExecuteAndProject_NegativeTimeout()
     {
         var builder = Mock.Of<IObjectBuilder<object>>();
 
         using var connection = new TestSqlConnection();
 
         Should.Throw<ArgumentOutOfRangeException>(() =>
         {
-            connection.ExecuteAndProjectTo("any", builder, timeout: -1);
+            connection.ExecuteAndProject("any", builder, timeout: -1);
         });
     }
 

PSql.Engine/ObjectResultSet.cs

@@ -12,6 +12,7 @@ internal sealed class ObjectResultSet<T> : IEnumerator<T>
     private readonly SqlDataReader     _reader;
     private readonly IObjectBuilder<T> _builder;
     private readonly bool              _useSqlTypes;
+    private readonly CancellationToken _cancellation;
 
     private T?        _current;
     private string[]? _columnNames;
@@ -20,16 +21,18 @@ public ObjectResultSet(
         SqlConnection     connection,
         SqlDataReader     reader,
         IObjectBuilder<T> builder,
-        bool              useSqlTypes)
+        bool              useSqlTypes,
+        CancellationToken cancellation = default)
     {
         ArgumentNullException.ThrowIfNull(connection);
         ArgumentNullException.ThrowIfNull(reader);
         ArgumentNullException.ThrowIfNull(builder);
 
-        _connection  = connection;
-        _reader      = reader;
-        _builder     = builder;
-        _useSqlTypes = useSqlTypes;
+        _connection   = connection;
+        _reader       = reader;
+        _builder      = builder;
+        _useSqlTypes  = useSqlTypes;
+        _cancellation = cancellation;
     }
 
     public T Current
@@ -40,9 +43,11 @@ public T Current
 
     public bool MoveNext()
     {
-        while (!_reader.Read())
+        // Using async from sync to gain cancellation support
+
+        while (!_reader.ReadAsync(_cancellation).GetAwaiter().GetResult())
         {
-            if (!_reader.NextResult())
+            if (!_reader.NextResultAsync(_cancellation).GetAwaiter().GetResult())
                 return SetNoCurrent();
 
             _columnNames = null;

PSql.Engine/SqlConnection.cs

@@ -167,10 +167,10 @@ protected SqlCommand SetUpCommand(
     /// <summary>
     ///   Ensures that the connection is open.
     /// </summary>
-    protected void AutoOpen()
+    protected void AutoOpen(CancellationToken cancellation)
     {
         if (Connection.State == ConnectionState.Closed)
-            Connection.Open();
+            Connection.OpenAsync(cancellation).GetAwaiter().GetResult();
     }
 
     /// <summary>
@@ -193,6 +193,9 @@ protected void AutoOpen()
     ///     the <see cref="System.Data.SqlTypes"/> namespace, such as
     ///     <see cref="System.Data.SqlTypes.SqlInt32"/>.
     /// </param>
+    /// <param name="cancellation">
+    ///   A token to monitor for cancellation requests.
+    /// </param>
     /// <returns>
     ///   A sequence of objects created by executing the <paramref name="sql"/>
     ///   batch and projecting each result row to an object using the
@@ -221,18 +224,21 @@ protected void AutoOpen()
     /// <exception cref="ObjectDisposedException">
     ///   Thrown by the underlying ADO.NET connection or command objects.
     /// </exception>
-    public IEnumerator<T> ExecuteAndProjectTo<T>(
+    public IEnumerator<T> ExecuteAndProject<T>(
         string            sql,
         IObjectBuilder<T> builder,
-        int               timeout     = 0,
-        bool              useSqlTypes = false)
+        int               timeout      = 0,
+        bool              useSqlTypes  = false,
+        CancellationToken cancellation = default)
     {
+        // Using async from sync to gain cancellation support
+
         SetUpCommand(sql, timeout);
-        AutoOpen();
+        AutoOpen(cancellation);
 
-        var reader = Command.ExecuteReader();
+        var reader = Command.ExecuteReaderAsync(cancellation).GetAwaiter().GetResult();
 
-        return new ObjectResultSet<T>(this, reader, builder, useSqlTypes);
+        return new ObjectResultSet<T>(this, reader, builder, useSqlTypes, cancellation);
     }
 
     /// <summary>

PSql/Commands/InvokeSqlCommand.cs

@@ -57,7 +57,8 @@ public class InvokeSqlCommand : ConnectedCmdlet
     [Parameter]
     public TimeSpan? Timeout { get; set; }
 
-    private readonly E.SqlCmdPreprocessor _preprocessor;
+    private readonly E.SqlCmdPreprocessor    _preprocessor;
+    private readonly CancellationTokenSource _cancellation;
 
     private bool ShouldUsePreprocessing
         => !NoPreprocessing;
@@ -68,6 +69,7 @@ private bool ShouldUseErrorHandling
     public InvokeSqlCommand()
     {
         _preprocessor = new();
+        _cancellation = new();
     }
 
     protected override void ProcessRecord()
@@ -90,6 +92,11 @@ protected override void ProcessRecord()
             Execute(scripts);
     }
 
+    protected override void StopProcessing()
+    {
+        _cancellation.Cancel();
+    }
+
     private static IEnumerable<string> ExcludeNullOrEmpty(IEnumerable<string?> scripts)
     {
         return scripts.Where(s => s.HasContent())!;
@@ -124,16 +131,16 @@ private IEnumerator<object> ExecuteAndProjectToObjects(string batch)
             ? (int) Timeout.Value.TotalSeconds
             : DefaultTimeoutSeconds;
 
-        return Connection.InnerConnection.ExecuteAndProjectTo(
-            batch, new PSObjectBuilder(), timeout, UseSqlTypes
+        return Connection.InnerConnection.ExecuteAndProject(
+            batch, new PSObjectBuilder(), timeout, UseSqlTypes, _cancellation.Token
         );
     }
 
-#if DECIDE_ON_CANCELLATION
-    protected override void StopProcessing()
+    /// <inheritdoc/>
+    public override void Dispose()
     {
-        // TODO: Figure out how to cancel
-        // Async from sync, perhaps?
+        _cancellation.Dispose();
+
+        base.Dispose();
     }
-#endif
 }