YdbCommand: Execute* methods don't propagate cancellation token to the server stream, there is used only for initializing the YdbDataReader.

Author: KirillKurdyukov

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,4 @@
+- `YdbCommand`: `Execute*` methods don't propagate cancellation token to the server stream, there is used only for initializing the `YdbDataReader`.
 - `YdbCommand`: Cancellation token propagation in `Execute*` methods
 - `YdbConnection`: Cancellation token propagation in `OpenAsync` method.
 - `YdbDataReader`: Cancellation token propagation in `ReadAsync` and `NextResultAsync` methods.

src/Ydb.Sdk/src/Ado/YdbCommand.cs

@@ -201,7 +201,6 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
         var execSettings = CommandTimeout > 0
             ? new ExecuteQuerySettings { TransportTimeout = TimeSpan.FromSeconds(CommandTimeout) }
             : new ExecuteQuerySettings();
-        execSettings.CancellationToken = cancellationToken;
 
         var transaction = YdbConnection.CurrentTransaction;
 
@@ -211,9 +210,9 @@ protected override async Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBeha
         }
 
         var ydbDataReader = await YdbDataReader.CreateYdbDataReader(
-            await YdbConnection.Session.ExecuteQuery(
-                preparedSql.ToString(), ydbParameters, execSettings, transaction?.TransactionControl
-            ), YdbConnection.OnStatus, transaction
+            await YdbConnection.Session
+                .ExecuteQuery(preparedSql.ToString(), ydbParameters, execSettings, transaction?.TransactionControl),
+            YdbConnection.OnStatus, transaction, cancellationToken
         );
 
         YdbConnection.LastReader = ydbDataReader;

src/Ydb.Sdk/src/Ado/YdbDataReader.cs

@@ -65,17 +65,19 @@ private YdbDataReader(
     internal static async Task<YdbDataReader> CreateYdbDataReader(
         IServerStream<ExecuteQueryResponsePart> resultSetStream,
         Action<Status> onStatus,
-        YdbTransaction? ydbTransaction = null)
+        YdbTransaction? ydbTransaction = null,
+        CancellationToken cancellationToken = default
+    )
     {
         var ydbDataReader = new YdbDataReader(resultSetStream, onStatus, ydbTransaction);
-        await ydbDataReader.Init();
+        await ydbDataReader.Init(cancellationToken);
 
         return ydbDataReader;
     }
 
-    private async Task Init()
+    private async Task Init(CancellationToken cancellationToken)
     {
-        if (State.IsConsumed == await NextExecPart(CancellationToken.None))
+        if (State.IsConsumed == await NextExecPart(cancellationToken))
         {
             throw new YdbException("YDB server closed the stream");
         }

src/Ydb.Sdk/tests/Ado/YdbConnectionTests.cs

@@ -261,6 +261,24 @@ public async Task YdbCommand_WhenCancelTokenIsCanceled_ThrowYdbException()
         Assert.True(ydbDataReader.IsClosed);
     }
 
+    [Fact]
+    public async Task ExecuteMethods_WhenExecutedYdbDataReaderThenCancelTokenIsCanceled_ReturnValues()
+    {
+        await using var connection = await CreateOpenConnectionAsync();
+        var ydbCommand = new YdbCommand(connection) { CommandText = "SELECT 1; SELECT 1; "};
+        var cts = new CancellationTokenSource();
+        var ydbDataReader = await ydbCommand.ExecuteReaderAsync(cts.Token);
+        
+        await ydbDataReader.ReadAsync(cts.Token);
+        Assert.Equal(1, ydbDataReader.GetValue(0));
+        Assert.True(await ydbDataReader.NextResultAsync(cts.Token));
+        cts.Cancel();
+        await ydbDataReader.ReadAsync(cts.Token);
+        Assert.Equal(1, ydbDataReader.GetValue(0));
+        // ReSharper disable once MethodSupportsCancellation
+        Assert.False(await ydbDataReader.NextResultAsync());
+    }
+
     private List<Task> GenerateTasks() => Enumerable.Range(0, 100).Select(async i =>
     {
         await using var connection = await CreateOpenConnectionAsync();