Use cancellation token in DbCommand.ExecuteReaderAsync call in DAB layer (#3300)

This is to address #3301

Use cancellation token in DbCommand.ExecuteReaderAsync call in DAB layer

Problem:
RequestTimeoutAttribute only works during graphql workload execution, it
doesn’t take effect during query execution as the logic to use
cancellation token to DAB layer is NOT implemented, even though the code
to pass cancellation token to DAB is already implemented.

Fix:
Use cancellation token to DAB layer in DbCommand.ExecuteReaderAsync call

## Why make this change?

- This is the 3rd task to address multiple GraphQL timeout related
issues. Discussion:
https://microsoft.sharepoint.com/:w:/r/teams/dbsaaspillar/_layouts/15/doc2.aspx?sourcedoc=%7B230dcf09-f556-4071-877a-8294b7df7338%7D&action=edit&wdPid=62b0c291&share=IQEJzw0jVvVxQId6gpS333M4AS2c9wqqYSgUlkTBDxQSCEw

## What is this change?

- Summary of how your changes work to give reviewers context of your
intent.
- Use cancellation token(if any) in DbCommand.ExecuteReaderAsync call in
DAB layer

## How was this tested?

- [ ] Integration Tests 
- [X] Unit Tests
- Also did manual test with long running(3 minutes) SQL query from
graphql and cancellation token with 30 seconds timeout, from debugger I
can see code goes to DbCommand.ExecuteReaderAsync call in DAB layer and
times out after 30 seconds.

---------

Co-authored-by: Aniruddh Munde <anmunde@microsoft.com>
Co-authored-by: Anusha Kolan <anushakolan10@gmail.com>

Author: 3 people

src/Core/Resolvers/QueryExecutor.cs

@@ -293,8 +293,11 @@ public virtual TConnection CreateConnection(string dataSourceName)
                 TResult? result = default(TResult);
                 try
                 {
-                    using DbDataReader dbDataReader = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ?
-                        await cmd.ExecuteReaderAsync(CommandBehavior.SequentialAccess) : await cmd.ExecuteReaderAsync(CommandBehavior.CloseConnection);
+                    CommandBehavior commandBehavior = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ? CommandBehavior.SequentialAccess : CommandBehavior.CloseConnection;
+                    // CancellationToken is passed to ExecuteReaderAsync to ensure that if the client times out while the query is executing, the execution will be cancelled and resources will be freed up.
+                    CancellationToken cancellationToken = httpContext?.RequestAborted ?? CancellationToken.None;
+                    using DbDataReader dbDataReader = await cmd.ExecuteReaderAsync(commandBehavior, cancellationToken);
+
                     if (dataReaderHandler is not null && dbDataReader is not null)
                     {
                         result = await dataReaderHandler(dbDataReader, args);

src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs

@@ -11,6 +11,7 @@
 using System.Linq;
 using System.Net;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
 using Azure.DataApiBuilder.Config;
@@ -1015,6 +1016,209 @@ private static Mock<IHttpContextAccessor> CreateHttpContextAccessorWithAuthentic
         #endregion
 
         /// <summary>
+        /// Validates that when the CancellationToken from httpContext.RequestAborted times out
+        /// during a long-running query execution (simulating ExecuteReaderAsync being interrupted
+        /// by a token timeout), the resulting TaskCanceledException propagates through the Polly
+        /// retry policy without any retry attempts.
+        /// Unlike TestCancellationExceptionIsNotRetriedByRetryPolicy which throws immediately,
+        /// this test simulates a real timeout where the cancellation occurs asynchronously
+        /// after a delay.
+        /// </summary>
+        [TestMethod, TestCategory(TestCategory.MSSQL)]
+        public async Task TestCancellationTokenTimeoutDuringQueryExecutionAsync()
+        {
+            RuntimeConfig mockConfig = new(
+               Schema: "",
+               DataSource: new(DatabaseType.MSSQL, "", new()),
+               Runtime: new(
+                   Rest: new(),
+                   GraphQL: new(),
+                   Mcp: new(),
+                   Host: new(null, null)
+               ),
+               Entities: new(new Dictionary<string, Entity>())
+           );
+
+            MockFileSystem fileSystem = new();
+            fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(mockConfig.ToJson()));
+            FileSystemRuntimeConfigLoader loader = new(fileSystem);
+            RuntimeConfigProvider provider = new(loader)
+            {
+                IsLateConfigured = true
+            };
+
+            Mock<ILogger<QueryExecutor<SqlConnection>>> queryExecutorLogger = new();
+            Mock<IHttpContextAccessor> httpContextAccessor = new();
+            HttpContext context = new DefaultHttpContext();
+            httpContextAccessor.Setup(x => x.HttpContext).Returns(context);
+            DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider);
+            Mock<MsSqlQueryExecutor> queryExecutor
+                = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null);
+
+            queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary<string, DbConnectionStringBuilder>());
+
+            queryExecutor.Setup(x => x.CreateConnection(
+               It.IsAny<string>())).CallBase();
+
+            // Set up a CancellationTokenSource that times out after a short delay,
+            // simulating httpContext.RequestAborted firing due to a client timeout.
+            CancellationTokenSource cts = new();
+            cts.CancelAfter(TimeSpan.FromMilliseconds(100));
+            context.RequestAborted = cts.Token;
+
+            // Mock ExecuteQueryAgainstDbAsync to simulate a long-running database query
+            // that is interrupted when the CancellationToken times out.
+            // Task.Delay with the cancellation token throws TaskCanceledException when the
+            // token fires, mimicking cmd.ExecuteReaderAsync being cancelled by a timed-out token.
+            // The Stopwatch + finally block mirrors the real ExecuteQueryAgainstDbAsync to verify
+            // that execution time is recorded even when a timeout occurs.
+            queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync(
+                It.IsAny<SqlConnection>(),
+                It.IsAny<string>(),
+                It.IsAny<IDictionary<string, DbConnectionParam>>(),
+                It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
+                It.IsAny<HttpContext>(),
+                provider.GetConfig().DefaultDataSourceName,
+                It.IsAny<List<string>>()))
+            .Returns(async () =>
+            {
+                Stopwatch timer = Stopwatch.StartNew();
+                try
+                {
+                    // Simulate a long-running query interrupted by token timeout.
+                    // Timeout.Infinite (-1) means "wait forever" — the only way this
+                    // completes is when cts.Token fires after ~100 ms, which causes
+                    // Task.Delay to throw TaskCanceledException.
+                    await Task.Delay(Timeout.Infinite, cts.Token);
+                    return (object)null;
+                }
+                finally
+                {
+                    timer.Stop();
+                    queryExecutor.Object.AddDbExecutionTimeToMiddlewareContext(timer.ElapsedMilliseconds);
+                }
+            });
+
+            // Call the actual ExecuteQueryAsync method (includes Polly retry policy).
+            queryExecutor.Setup(x => x.ExecuteQueryAsync(
+                It.IsAny<string>(),
+                It.IsAny<IDictionary<string, DbConnectionParam>>(),
+                It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
+                It.IsAny<string>(),
+                It.IsAny<HttpContext>(),
+                It.IsAny<List<string>>())).CallBase();
+
+            // Act & Assert: TaskCanceledException should propagate without retries.
+            await Assert.ThrowsExceptionAsync<TaskCanceledException>(async () =>
+            {
+                await queryExecutor.Object.ExecuteQueryAsync<object>(
+                    sqltext: string.Empty,
+                    parameters: new Dictionary<string, DbConnectionParam>(),
+                    dataReaderHandler: null,
+                    dataSourceName: String.Empty,
+                    httpContext: context,
+                    args: null);
+            });
+
+            // Verify that the underlying database execution is invoked exactly once,
+            // confirming that Polly does not perform any retries for TaskCanceledException.
+            queryExecutor.Verify(q => q.ExecuteQueryAgainstDbAsync(
+                    It.IsAny<SqlConnection>(),
+                    It.IsAny<string>(),
+                    It.IsAny<IDictionary<string, DbConnectionParam>>(),
+                    It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
+                    It.IsAny<HttpContext>(),
+                    provider.GetConfig().DefaultDataSourceName,
+                    It.IsAny<List<string>>()),
+                Times.Once);
+
+            // Verify the finally block recorded execution time even though the token timed out.
+            Assert.IsTrue(
+                context.Items.ContainsKey(TOTAL_DB_EXECUTION_TIME),
+                "HttpContext must contain the total db execution time even when the request is cancelled.");
+        }
+
+        /// <summary>
+        /// Validates that when ExecuteQueryAgainstDbAsync throws OperationCanceledException
+        /// (e.g., due to client disconnect via httpContext.RequestAborted cancellation token),
+        /// the Polly retry policy does NOT retry and the exception propagates to the caller.
+        /// The retry policy is configured to only handle DbException, so OperationCanceledException
+        /// should be immediately re-thrown without any retry attempts.
+        /// </summary>
+        [TestMethod, TestCategory(TestCategory.MSSQL)]
+        public async Task TestCancellationExceptionIsNotRetriedByRetryPolicy()
+        {
+            RuntimeConfig mockConfig = new(
+               Schema: "",
+               DataSource: new(DatabaseType.MSSQL, "", new()),
+               Runtime: new(
+                   Rest: new(),
+                   GraphQL: new(),
+                   Mcp: new(),
+                   Host: new(null, null)
+               ),
+               Entities: new(new Dictionary<string, Entity>())
+           );
+
+            MockFileSystem fileSystem = new();
+            fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(mockConfig.ToJson()));
+            FileSystemRuntimeConfigLoader loader = new(fileSystem);
+            RuntimeConfigProvider provider = new(loader)
+            {
+                IsLateConfigured = true
+            };
+
+            Mock<ILogger<QueryExecutor<SqlConnection>>> queryExecutorLogger = new();
+            Mock<IHttpContextAccessor> httpContextAccessor = new();
+            DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider);
+            Mock<MsSqlQueryExecutor> queryExecutor
+                = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null);
+
+            queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary<string, DbConnectionStringBuilder>());
+
+            queryExecutor.Setup(x => x.CreateConnection(
+               It.IsAny<string>())).CallBase();
+
+            // Mock ExecuteQueryAgainstDbAsync to throw OperationCanceledException,
+            // simulating a cancelled CancellationToken from httpContext.RequestAborted.
+            queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync(
+                It.IsAny<SqlConnection>(),
+                It.IsAny<string>(),
+                It.IsAny<IDictionary<string, DbConnectionParam>>(),
+                It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
+                It.IsAny<HttpContext>(),
+                provider.GetConfig().DefaultDataSourceName,
+                It.IsAny<List<string>>()))
+            .ThrowsAsync(new OperationCanceledException("The operation was canceled."));
+
+            // Call the actual ExecuteQueryAsync method.
+            queryExecutor.Setup(x => x.ExecuteQueryAsync(
+                It.IsAny<string>(),
+                It.IsAny<IDictionary<string, DbConnectionParam>>(),
+                It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
+                It.IsAny<string>(),
+                It.IsAny<HttpContext>(),
+                It.IsAny<List<string>>())).CallBase();
+
+            // Act & Assert: OperationCanceledException should propagate without retries.
+            await Assert.ThrowsExceptionAsync<OperationCanceledException>(async () =>
+            {
+                await queryExecutor.Object.ExecuteQueryAsync<object>(
+                    sqltext: string.Empty,
+                    parameters: new Dictionary<string, DbConnectionParam>(),
+                    dataReaderHandler: null,
+                    dataSourceName: String.Empty,
+                    httpContext: null,
+                    args: null);
+            });
+
+            // Verify no retry log messages were emitted. Since IsLateConfigured is true,
+            // the debug log is skipped, and since Polly doesn't handle OperationCanceledException,
+            // no retry occurs → zero logger invocations.
+            Assert.AreEqual(0, queryExecutorLogger.Invocations.Count);
+        }
+
+        /// <summary>  
         /// Validates that GetSessionParamsQuery includes all observability values:
         /// - OpenTelemetry correlation values (dab.trace_id, dab.span_id) when an Activity is present
         /// - OBO observability values (dab.auth_type, dab.user_id, dab.tenant_id) when user-delegated auth is enabled