diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 97e2f7e8d4..2c1d080bef 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); + var 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 778144c6d2..9ec1e536e3 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; @@ -1014,6 +1015,202 @@ private static Mock CreateHttpContextAccessorWithAuthentic #endregion + /// + /// 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 no retry log messages were emitted. Polly does not handle + // TaskCanceledException (subclass of OperationCanceledException), so + // the exception propagates immediately without any retry attempts. + Assert.AreEqual(0, queryExecutorLogger.Invocations.Count); + + // 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() {