feat: add support for CancellationToken in generated methods

.github/workflows/legacy-tests.yml

@@ -70,8 +70,11 @@ jobs:
           timeout_minutes: 7
           shell: pwsh
           command: |
-            $mysqlJob = Start-Job -ScriptBlock { 
-              choco install mysql --no-progress --version=8.4.6 -y --params "/serviceName:MySQL" 
+            $mysqlJob = Start-Job -ScriptBlock {
+              # choco has no working 8.4.x mysql package (8.4.6's zip was pulled from cdn.mysql.com,
+              # no 8.4.7+ package exists), so we track the latest available. If this 404s again,
+              # switch to a dedicated setup action (shogo82148/actions-setup-mysql) for version pinning.
+              choco install mysql --no-progress --version=9.6.0 -y --params "/serviceName:MySQL"
               return $LASTEXITCODE
             }
 
@@ -114,4 +117,4 @@ jobs:
         run: |
           $path = vswhere -latest -products * -requires Microsoft.VisualStudio.Workload.ManagedDesktop Microsoft.VisualStudio.Workload.Web -requiresAny -property installationPath
           $path = join-path $path 'Common7\IDE\CommonExtensions\Microsoft\TestWindow\vstest.console.exe'
-          & $path ./end2end/EndToEndTestsLegacy/bin/Release/net472/EndToEndTestsLegacy.dll
+          & $path ./end2end/EndToEndTestsLegacy/bin/Release/net472/EndToEndTestsLegacy.dll

Drivers/DbDriver.cs

@@ -1,5 +1,6 @@
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Plugin;
+using SqlcGenCsharp.Drivers.Generators;
 using System;
 using System.Collections.Generic;
 using System.Linq;
@@ -42,6 +43,8 @@ public abstract class DbDriver
 
     public Options Options { get; }
 
+    public CancellationGen Cancellation => new(Options.WithCancellationToken);
+
     public string DefaultSchema { get; }
 
     public abstract string TransactionClassName { get; }
@@ -296,7 +299,7 @@ public virtual string[] GetLastIdStatement(Query query)
         var convertFuncCall = convertFunc(Variable.Result.AsVarName());
         return
         [
-            $"var {Variable.Result.AsVarName()} = await {Variable.Command.AsVarName()}.ExecuteScalarAsync();",
+            $"var {Variable.Result.AsVarName()} = await {Variable.Command.AsVarName()}.ExecuteScalarAsync({Cancellation.Argument()});",
             $"return {convertFuncCall};"
         ];
     }

Drivers/Generators/CancellationGen.cs

@@ -0,0 +1,38 @@
+namespace SqlcGenCsharp.Drivers.Generators;
+
+// Owns every way an optional CancellationToken shows up in generated code.
+// When disabled, every fragment is empty / a pass-through, so output is unchanged.
+public class CancellationGen(bool enabled)
+{
+    private static string TokenName => Variable.CancellationToken.AsVarName();
+
+    // Trailing optional parameter for a generated method signature.
+    public string MethodParameter()
+    {
+        return enabled ? $"CancellationToken {TokenName} = default" : string.Empty;
+    }
+
+    // Trailing optional parameter appended after an existing parameter list, e.g. (List<T> args{TrailingMethodParameter()}).
+    public string TrailingMethodParameter()
+    {
+        return enabled ? $", CancellationToken {TokenName} = default" : string.Empty;
+    }
+
+    // Sole argument for an otherwise parameterless async call, e.g. ReadAsync(Argument()).
+    public string Argument()
+    {
+        return enabled ? TokenName : string.Empty;
+    }
+
+    // Trailing argument appended after existing call arguments, e.g. WriteAsync(value{TrailingArgument()}).
+    public string TrailingArgument()
+    {
+        return enabled ? $", {TokenName}" : string.Empty;
+    }
+
+    // Wraps flat Dapper call arguments in a CommandDefinition carrying the token; pass-through when disabled.
+    public string WrapDapperArgs(string flatArgs)
+    {
+        return enabled ? $"new CommandDefinition({flatArgs}, cancellationToken: {TokenName})" : flatArgs;
+    }
+}

Drivers/Generators/CommonGen.cs

@@ -6,11 +6,13 @@ namespace SqlcGenCsharp.Drivers.Generators;
 
 public class CommonGen(DbDriver dbDriver)
 {
-    public static string GetMethodParameterList(string argInterface, IEnumerable<Parameter> parameters)
+    public static string GetMethodParameterList(string argInterface, IEnumerable<Parameter> parameters,
+        string cancellationParam = "")
     {
-        return $"{(string.IsNullOrEmpty(argInterface) || !parameters.Any()
+        var argsParam = string.IsNullOrEmpty(argInterface) || !parameters.Any()
             ? string.Empty
-            : $"{argInterface} {Variable.Args.AsVarName()}")}";
+            : $"{argInterface} {Variable.Args.AsVarName()}";
+        return string.Join(", ", new[] { argsParam, cancellationParam }.Where(s => !string.IsNullOrEmpty(s)));
     }
 
     public static string GetDapperArgs(Query query)
@@ -52,14 +54,14 @@ public string ConstructDapperParamsDict(Query query)
                  """;
     }
 
-    public static string AwaitReaderRow()
+    public static string AwaitReaderRow(string cancellationArg)
     {
-        return $"await {Variable.Reader.AsVarName()}.ReadAsync()";
+        return $"await {Variable.Reader.AsVarName()}.ReadAsync({cancellationArg})";
     }
 
-    public static string InitDataReader()
+    public static string InitDataReader(string cancellationArg)
     {
-        return $"var {Variable.Reader.AsVarName()} = await {Variable.Command.AsVarName()}.ExecuteReaderAsync()";
+        return $"var {Variable.Reader.AsVarName()} = await {Variable.Command.AsVarName()}.ExecuteReaderAsync({cancellationArg})";
     }
 
     public static string GetSqlTransformations(Query query, string queryTextConstant)

Drivers/Generators/CopyFromDeclareGen.cs

@@ -9,7 +9,7 @@ public class CopyFromDeclareGen(DbDriver dbDriver)
     public MemberDeclarationSyntax Generate(string queryTextConstant, string argInterface, Query query)
     {
         return ParseMemberDeclaration($$"""
-            public async Task {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}(List<{{argInterface}}> args)
+            public async Task {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}(List<{{argInterface}}> args{{dbDriver.Cancellation.TrailingMethodParameter()}})
             {
                 {{((ICopyFrom)dbDriver).GetCopyFromImpl(query, queryTextConstant)}}
             }

Drivers/Generators/ExecDeclareGen.cs

@@ -10,7 +10,7 @@ public class ExecDeclareGen(DbDriver dbDriver)
 
     public MemberDeclarationSyntax Generate(string queryTextConstant, string argInterface, Query query)
     {
-        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params);
+        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params, dbDriver.Cancellation.MethodParameter());
         return ParseMemberDeclaration($$"""
             public async Task {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}({{parametersStr}})
             {
@@ -46,9 +46,10 @@ private string GetDapperNoTxBody(string sqlVar, Query query)
     {
         var connectionCommands = dbDriver.EstablishConnection(query);
         var dapperArgs = CommonGen.GetDapperArgs(query);
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}");
         return connectionCommands.GetConnectionOrDataSource.WrapBlock(
             $"""
-            await {Variable.Connection.AsVarName()}.ExecuteAsync({sqlVar}{dapperArgs});
+            await {Variable.Connection.AsVarName()}.ExecuteAsync({callArgs});
             return;
             """
         );
@@ -58,6 +59,15 @@ private string GetDapperWithTxBody(string sqlVar, Query query)
     {
         var transactionProperty = Variable.Transaction.AsPropertyName();
         var dapperArgs = CommonGen.GetDapperArgs(query);
+        if (dbDriver.Options.WithCancellationToken)
+        {
+            var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
+            return $$"""
+                        {{dbDriver.TransactionConnectionNullExcetionThrow}}
+                        await this.{{transactionProperty}}.Connection.ExecuteAsync({{callArgs}});
+                     """;
+        }
+
         return $$"""
                     {{dbDriver.TransactionConnectionNullExcetionThrow}}
                     await this.{{transactionProperty}}.Connection.ExecuteAsync(
@@ -75,7 +85,7 @@ private string GetDriverNoTxBody(string sqlVar, Query query)
             {sqlCommands.SetCommandText.AppendSemicolonUnlessEmpty()}
             {dbDriver.AddParametersToCommand(query)}
             {sqlCommands.PrepareCommand.AppendSemicolonUnlessEmpty()}
-            await {Variable.Command.AsVarName()}.ExecuteNonQueryAsync();
+            await {Variable.Command.AsVarName()}.ExecuteNonQueryAsync({dbDriver.Cancellation.Argument()});
             """
         );
         return connectionCommands.GetConnectionOrDataSource.WrapBlock(
@@ -99,7 +109,7 @@ private string GetDriverWithTxBody(string sqlVar, Query query)
                 {{commandVar}}.CommandText = {{sqlVar}};
                 {{commandVar}}.Transaction = this.{{transactionProperty}};
                 {{dbDriver.AddParametersToCommand(query)}}
-                await {{commandVar}}.ExecuteNonQueryAsync();
+                await {{commandVar}}.ExecuteNonQueryAsync({{dbDriver.Cancellation.Argument()}});
             }
         """;
     }

Drivers/Generators/ExecLastIdDeclareGen.cs

@@ -11,7 +11,7 @@ public class ExecLastIdDeclareGen(DbDriver dbDriver)
 
     public MemberDeclarationSyntax Generate(string queryTextConstant, string argInterface, Query query)
     {
-        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params);
+        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params, dbDriver.Cancellation.MethodParameter());
         return ParseMemberDeclaration($$"""
             public async Task<{{dbDriver.GetIdColumnType(query)}}> {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}({{parametersStr}})
             {
@@ -48,18 +48,20 @@ private string GetDapperNoTxBody(string sqlVar, Query query)
         var connectionCommands = dbDriver.EstablishConnection(query);
         var dapperArgs = CommonGen.GetDapperArgs(query);
         var idColumnType = dbDriver.GetIdColumnType(query);
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}");
         return connectionCommands.GetConnectionOrDataSource.WrapBlock($"""
-            return await {Variable.Connection.AsVarName()}.QuerySingleAsync<{idColumnType}>({sqlVar}{dapperArgs});
+            return await {Variable.Connection.AsVarName()}.QuerySingleAsync<{idColumnType}>({callArgs});
         """);
     }
 
     private string GetDapperWithTxBody(string sqlVar, Query query)
     {
         var transactionProperty = Variable.Transaction.AsPropertyName();
         var dapperArgs = query.Params.Any() ? $", {Variable.QueryParams.AsVarName()}" : string.Empty;
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
         return $$"""
                 {{dbDriver.TransactionConnectionNullExcetionThrow}}
-                return await this.{{transactionProperty}}.Connection.QuerySingleAsync<{{dbDriver.GetIdColumnType(query)}}>({{sqlVar}}{{dapperArgs}}, transaction: this.{{transactionProperty}});
+                return await this.{{transactionProperty}}.Connection.QuerySingleAsync<{{dbDriver.GetIdColumnType(query)}}>({{callArgs}});
                 """;
     }
 

Drivers/Generators/ExecRowsDeclareGen.cs

@@ -10,7 +10,7 @@ public class ExecRowsDeclareGen(DbDriver dbDriver)
 
     public MemberDeclarationSyntax Generate(string queryTextConstant, string argInterface, Query query)
     {
-        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params);
+        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params, dbDriver.Cancellation.MethodParameter());
         return ParseMemberDeclaration($$"""
             public async Task<long> {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}({{parametersStr}})
             {
@@ -46,15 +46,25 @@ private string GetDapperNoTxBody(string sqlVar, Query query)
     {
         var connectionCommands = dbDriver.EstablishConnection(query);
         var dapperArgs = CommonGen.GetDapperArgs(query);
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}");
         return connectionCommands.GetConnectionOrDataSource.WrapBlock(
-            $"return await {Variable.Connection.AsVarName()}.ExecuteAsync({sqlVar}{dapperArgs});"
+            $"return await {Variable.Connection.AsVarName()}.ExecuteAsync({callArgs});"
         );
     }
 
     private string GetDapperWithTxBody(string sqlVar, Query query)
     {
         var transactionProperty = Variable.Transaction.AsPropertyName();
         var dapperArgs = CommonGen.GetDapperArgs(query);
+        if (dbDriver.Options.WithCancellationToken)
+        {
+            var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
+            return $$"""
+                {{dbDriver.TransactionConnectionNullExcetionThrow}}
+                return await this.{{transactionProperty}}.Connection.ExecuteAsync({{callArgs}});
+            """;
+        }
+
         return $$"""
             {{dbDriver.TransactionConnectionNullExcetionThrow}}
             return await this.{{transactionProperty}}.Connection.ExecuteAsync(
@@ -72,7 +82,7 @@ private string GetDriverNoTxBody(string sqlVar, Query query)
             {sqlCommands.SetCommandText.AppendSemicolonUnlessEmpty()}
             {dbDriver.AddParametersToCommand(query)}
             {sqlCommands.PrepareCommand.AppendSemicolonUnlessEmpty()}
-            return await {Variable.Command.AsVarName()}.ExecuteNonQueryAsync();
+            return await {Variable.Command.AsVarName()}.ExecuteNonQueryAsync({dbDriver.Cancellation.Argument()});
             """
         );
         return connectionCommands.GetConnectionOrDataSource.WrapBlock(
@@ -95,7 +105,7 @@ private string GetDriverWithTxBody(string sqlVar, Query query)
                 {{commandVar}}.CommandText = {{sqlVar}};
                 {{commandVar}}.Transaction = this.{{transactionProperty}};
                 {{dbDriver.AddParametersToCommand(query)}}
-                return await {{commandVar}}.ExecuteNonQueryAsync();
+                return await {{commandVar}}.ExecuteNonQueryAsync({{dbDriver.Cancellation.Argument()}});
             }
         """;
     }

Drivers/Generators/ManyDeclareGen.cs

@@ -12,7 +12,7 @@ public class ManyDeclareGen(DbDriver dbDriver)
 
     public MemberDeclarationSyntax Generate(string queryTextConstant, string argInterface, string returnInterface, Query query)
     {
-        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params);
+        var parametersStr = CommonGen.GetMethodParameterList(argInterface, query.Params, dbDriver.Cancellation.MethodParameter());
         var returnType = $"Task<List<{returnInterface}>>";
         return ParseMemberDeclaration($$"""
             public async {{returnType}} {{query.Name.ToMethodName(dbDriver.Options.WithAsyncSuffix)}}({{parametersStr}})
@@ -52,9 +52,10 @@ private string GetDapperNoTxBody(string sqlVar, string returnInterface, Query qu
         var dapperArgs = CommonGen.GetDapperArgs(query);
         var returnType = dbDriver.AddNullableSuffixIfNeeded(returnInterface, true);
         var resultVar = Variable.Result.AsVarName();
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}");
         return connectionCommands.GetConnectionOrDataSource.WrapBlock(
             $"""
-            var {resultVar} = await {Variable.Connection.AsVarName()}.QueryAsync<{returnType}>({sqlVar}{dapperArgs});
+            var {resultVar} = await {Variable.Connection.AsVarName()}.QueryAsync<{returnType}>({callArgs});
             return {resultVar}.AsList();
             """
         );
@@ -66,6 +67,15 @@ private string GetDapperWithTxBody(string sqlVar, string returnInterface, Query
         var dapperArgs = CommonGen.GetDapperArgs(query);
         var returnType = dbDriver.AddNullableSuffixIfNeeded(returnInterface, true);
 
+        if (dbDriver.Options.WithCancellationToken)
+        {
+            var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
+            return $$"""
+                {{dbDriver.TransactionConnectionNullExcetionThrow}}
+                return (await this.{{transactionProperty}}.Connection.QueryAsync<{{returnType}}>({{callArgs}})).AsList();
+            """;
+        }
+
         return $$"""
             {{dbDriver.TransactionConnectionNullExcetionThrow}}
             return (await this.{{transactionProperty}}.Connection.QueryAsync<{{returnType}}>(
@@ -80,7 +90,7 @@ private string GetDriverNoTxBody(string sqlVar, string returnInterface, Query qu
         var dataclassInit = CommonGen.InstantiateDataclass([.. query.Columns], returnInterface, query);
         var resultVar = Variable.Result.AsVarName();
         var readWhileExists = $$"""
-            while ({{CommonGen.AwaitReaderRow()}})
+            while ({{CommonGen.AwaitReaderRow(dbDriver.Cancellation.Argument())}})
                 {{resultVar}}.Add({{dataclassInit}});
         """;
         var sqlCommands = dbDriver.CreateSqlCommand(sqlVar);
@@ -89,7 +99,7 @@ private string GetDriverNoTxBody(string sqlVar, string returnInterface, Query qu
             {{sqlCommands.SetCommandText.AppendSemicolonUnlessEmpty()}}
             {{dbDriver.AddParametersToCommand(query)}}
             {{sqlCommands.PrepareCommand.AppendSemicolonUnlessEmpty()}}
-            using ({{CommonGen.InitDataReader()}})
+            using ({{CommonGen.InitDataReader(dbDriver.Cancellation.Argument())}})
             {
                 var {{resultVar}} = new List<{{returnInterface}}>();
                 {{readWhileExists}}
@@ -118,10 +128,10 @@ private string GetDriverWithTxBody(string sqlVar, string returnInterface, Query
                 {{commandVar}}.CommandText = {{sqlVar}};
                 {{commandVar}}.Transaction = this.{{transactionProperty}};
                 {{dbDriver.AddParametersToCommand(query)}}
-                using ({{CommonGen.InitDataReader()}})
+                using ({{CommonGen.InitDataReader(dbDriver.Cancellation.Argument())}})
                 {
                     var {{resultVar}} = new List<{{returnInterface}}>();
-                    while ({{CommonGen.AwaitReaderRow()}})
+                    while ({{CommonGen.AwaitReaderRow(dbDriver.Cancellation.Argument())}})
                         {{resultVar}}.Add({{CommonGen.InstantiateDataclass([.. query.Columns], returnInterface, query)}});
                     return {{resultVar}};
                 }