From 5747b13a0ead50de1ccc1e76639827145be135ff Mon Sep 17 00:00:00 2001 From: naxing123 <53497183+naxing123@users.noreply.github.com> Date: Mon, 30 Mar 2026 10:04:37 -0700 Subject: [PATCH 1/2] Use cancellation token in DbCommand.ExecuteReaderAsync call in DAB layer (#3300) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is to address https://github.com/Azure/data-api-builder/issues/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 Co-authored-by: Anusha Kolan --- src/Core/Resolvers/QueryExecutor.cs | 7 +- .../UnitTests/SqlQueryExecutorUnitTests.cs | 204 ++++++++++++++++++ 2 files changed, 209 insertions(+), 2 deletions(-) diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 97e2f7e8d4..d96b19af38 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/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); diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index b3782950f9..6cb0ef4cb6 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/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; @@ -669,6 +670,209 @@ public void ValidateStreamingLogicForEmptyCellsAsync() Assert.AreEqual(availableSize, (int)runtimeConfig.MaxResponseSizeMB() * 1024 * 1024); } + /// + /// 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. + /// + [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()) + ); + + 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>> queryExecutorLogger = new(); + Mock httpContextAccessor = new(); + HttpContext context = new DefaultHttpContext(); + httpContextAccessor.Setup(x => x.HttpContext).Returns(context); + DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); + Mock queryExecutor + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); + + queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); + + queryExecutor.Setup(x => x.CreateConnection( + It.IsAny())).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(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>())) + .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(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>())).CallBase(); + + // Act & Assert: TaskCanceledException should propagate without retries. + await Assert.ThrowsExceptionAsync(async () => + { + await queryExecutor.Object.ExecuteQueryAsync( + sqltext: string.Empty, + parameters: new Dictionary(), + 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(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>()), + 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."); + } + + /// + /// 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. + /// + [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()) + ); + + 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>> queryExecutorLogger = new(); + Mock httpContextAccessor = new(); + DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); + Mock queryExecutor + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); + + queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); + + queryExecutor.Setup(x => x.CreateConnection( + It.IsAny())).CallBase(); + + // Mock ExecuteQueryAgainstDbAsync to throw OperationCanceledException, + // simulating a cancelled CancellationToken from httpContext.RequestAborted. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>())) + .ThrowsAsync(new OperationCanceledException("The operation was canceled.")); + + // Call the actual ExecuteQueryAsync method. + queryExecutor.Setup(x => x.ExecuteQueryAsync( + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>())).CallBase(); + + // Act & Assert: OperationCanceledException should propagate without retries. + await Assert.ThrowsExceptionAsync(async () => + { + await queryExecutor.Object.ExecuteQueryAsync( + sqltext: string.Empty, + parameters: new Dictionary(), + 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); + } + [TestCleanup] public void CleanupAfterEachTest() { From d7458f42808284405d6ea2d2d9dc1e27ba5d0ccf Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Wed, 1 Apr 2026 15:18:33 -0700 Subject: [PATCH 2/2] fix unit tests due to some unit test class change only in main but not release branch --- src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index 6cb0ef4cb6..fb0898a642 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -708,7 +708,7 @@ public async Task TestCancellationTokenTimeoutDuringQueryExecutionAsync() httpContextAccessor.Setup(x => x.HttpContext).Returns(context); DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); Mock queryExecutor - = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null); queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); @@ -827,7 +827,7 @@ public async Task TestCancellationExceptionIsNotRetriedByRetryPolicy() Mock httpContextAccessor = new(); DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); Mock queryExecutor - = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null); queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary());