feat: add support for CancellationToken in generated methods

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,11 +59,10 @@ private string GetDapperWithTxBody(string sqlVar, Query query)
     {
         var transactionProperty = Variable.Transaction.AsPropertyName();
         var dapperArgs = CommonGen.GetDapperArgs(query);
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
         return $$"""
                     {{dbDriver.TransactionConnectionNullExcetionThrow}}
-                    await this.{{transactionProperty}}.Connection.ExecuteAsync(
-                            {{sqlVar}}{{dapperArgs}},
-                            transaction: this.{{transactionProperty}});
+                    await this.{{transactionProperty}}.Connection.ExecuteAsync({{callArgs}});
                  """;
     }
 
@@ -75,7 +75,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 +99,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,20 +46,20 @@ 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);
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
         return $$"""
             {{dbDriver.TransactionConnectionNullExcetionThrow}}
-            return await this.{{transactionProperty}}.Connection.ExecuteAsync(
-                    {{sqlVar}}{{dapperArgs}},
-                    transaction: this.{{transactionProperty}});
+            return await this.{{transactionProperty}}.Connection.ExecuteAsync({{callArgs}});
         """;
     }
 
@@ -72,7 +72,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 +95,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,11 +67,10 @@ private string GetDapperWithTxBody(string sqlVar, string returnInterface, Query
         var dapperArgs = CommonGen.GetDapperArgs(query);
         var returnType = dbDriver.AddNullableSuffixIfNeeded(returnInterface, true);
 
+        var callArgs = dbDriver.Cancellation.WrapDapperArgs($"{sqlVar}{dapperArgs}, transaction: this.{transactionProperty}");
         return $$"""
             {{dbDriver.TransactionConnectionNullExcetionThrow}}
-            return (await this.{{transactionProperty}}.Connection.QueryAsync<{{returnType}}>(
-                    {{sqlVar}}{{dapperArgs}},
-                    transaction: this.{{transactionProperty}})).AsList();
+            return (await this.{{transactionProperty}}.Connection.QueryAsync<{{returnType}}>({{callArgs}})).AsList();
         """;
     }
 
@@ -80,7 +80,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 +89,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 +118,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}};
                 }