From 3ec5b55b778d09f383a848532acd81961042e4ad Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Thu, 19 Mar 2026 05:16:19 -0700 Subject: [PATCH 1/5] =?UTF-8?q?Task=202004135:=20RequestTimeoutAttribute?= =?UTF-8?q?=20won=E2=80=99t=20take=20effect=20during=20query=20execution?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit https://dev.azure.com/powerbi/AppDev/_workitems/edit/2004135 RequestTimeoutAttribute won’t take effect during query execution Problem: RequestTimeoutAttribute only works during graphql workload execution, it doesn’t take effect during query execution due to request executor does NOT implement cancellation token logic. Fix: pass cancellation token to DAB layer to DbCommand.ExecuteReaderAsync call --- src/Core/Resolvers/QueryExecutor.cs | 7 +- .../UnitTests/SqlQueryExecutorUnitTests.cs | 333 ++++++++++++++++++ 2 files changed, 338 insertions(+), 2 deletions(-) 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 b3782950f9..6c4076418e 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,338 @@ 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 TestCancellationTokenTimeoutDuringQueryExecution() + { + 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); + + 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. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>())) + .Returns(async () => + { + // 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; + }); + + // 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); + } + + /// + /// 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); + + 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); + } + + /// + /// Validates that when ExecuteQueryAgainstDbAsync throws TaskCanceledException + /// (a subclass of OperationCanceledException, commonly thrown when an HTTP request + /// is aborted), the retry policy does NOT retry the operation. + /// + [TestMethod, TestCategory(TestCategory.MSSQL)] + public async Task TestTaskCancelledExceptionIsNotRetriedByRetryPolicy() + { + 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); + + queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); + + queryExecutor.Setup(x => x.CreateConnection( + It.IsAny())).CallBase(); + + // Mock ExecuteQueryAgainstDbAsync to throw TaskCanceledException, + // simulating an HTTP request timeout or client disconnect. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>())) + .ThrowsAsync(new TaskCanceledException("A task was canceled.")); + + 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: null, + args: null); + }); + + // No retries should have occurred. + Assert.AreEqual(0, queryExecutorLogger.Invocations.Count); + } + + /// + /// Validates that when a request is cancelled during query execution, + /// the database execution time is still recorded in the HttpContext. + /// This ensures the finally block in ExecuteQueryAgainstDbAsync runs + /// even when the CancellationToken triggers an exception. + /// + [TestMethod, TestCategory(TestCategory.MSSQL)] + public async Task TestDbExecutionTimeRecordedEvenWhenRequestCancelled() + { + 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); + + queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); + + // Call the actual ExecuteQueryAgainstDbAsync to exercise the finally block. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>())).CallBase(); + + // Create an HttpContext with a pre-cancelled CancellationToken. + CancellationTokenSource cts = new(); + cts.Cancel(); + context.RequestAborted = cts.Token; + + try + { + await queryExecutor.Object.ExecuteQueryAgainstDbAsync( + conn: null, + sqltext: string.Empty, + parameters: new Dictionary(), + dataReaderHandler: null, + dataSourceName: String.Empty, + httpContext: context, + args: null); + } + catch (Exception) + { + // SqlConnection is sealed and can't be mocked; conn is null so OpenAsync + // throws NullReferenceException. Ignore to verify the finally block behavior. + } + + // The finally block should have recorded execution time even though an exception occurred. + Assert.IsTrue( + context.Items.ContainsKey(TOTAL_DB_EXECUTION_TIME), + "HttpContext must contain the total db execution time even when the request is cancelled."); + } + [TestCleanup] public void CleanupAfterEachTest() { From 8242a85381a6b6d14b8ecc2d6d65a4e3c038fda1 Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Thu, 19 Mar 2026 13:24:07 -0700 Subject: [PATCH 2/5] address some comments for test --- .../UnitTests/SqlQueryExecutorUnitTests.cs | 257 ++++++++++++++++++ 1 file changed, 257 insertions(+) diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index 8327caaf08..f9d1e1ba51 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1015,6 +1015,263 @@ 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); + + 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. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + provider.GetConfig().DefaultDataSourceName, + It.IsAny>())) + .Returns(async () => + { + // 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; + }); + + // 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); + } + + /// + /// 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); + + 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); + } + + /// + /// Validates that when a request is cancelled during query execution, + /// the database execution time is still recorded in the HttpContext. + /// This ensures the finally block in ExecuteQueryAgainstDbAsync runs + /// even when the CancellationToken triggers an exception. + /// + [TestMethod, TestCategory(TestCategory.MSSQL)] + public async Task TestDbExecutionTimeRecordedEvenWhenRequestCancelled() + { + 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); + + queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); + + // Call the actual ExecuteQueryAgainstDbAsync to exercise the finally block. + queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>(), + It.IsAny, Task>>(), + It.IsAny(), + It.IsAny(), + It.IsAny>())).CallBase(); + + // Create an HttpContext with a pre-cancelled CancellationToken. + CancellationTokenSource cts = new(); + cts.Cancel(); + context.RequestAborted = cts.Token; + + try + { + await queryExecutor.Object.ExecuteQueryAgainstDbAsync( + conn: null, + sqltext: string.Empty, + parameters: new Dictionary(), + dataReaderHandler: null, + dataSourceName: String.Empty, + httpContext: context, + args: null); + } + catch (Exception) + { + // SqlConnection is sealed and can't be mocked; conn is null so OpenAsync + // throws NullReferenceException. Ignore to verify the finally block behavior. + } + + // The finally block should have recorded execution time even though an exception occurred. + Assert.IsTrue( + context.Items.ContainsKey(TOTAL_DB_EXECUTION_TIME), + "HttpContext must contain the total db execution time even when the request is cancelled."); + } + [TestCleanup] public void CleanupAfterEachTest() { From c9aba8181fada9620a9eccaf7b783407f9442914 Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Thu, 19 Mar 2026 18:25:45 -0700 Subject: [PATCH 3/5] address PR comment and test failures: 1) merged two tests 2) updated the queryExecutor contrtructor parameter after merging main --- .../UnitTests/SqlQueryExecutorUnitTests.cs | 109 ++++-------------- 1 file changed, 24 insertions(+), 85 deletions(-) diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index f9d1e1ba51..9ec1e536e3 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1053,7 +1053,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); + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); @@ -1070,6 +1070,8 @@ Mock queryExecutor // 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(), @@ -1080,12 +1082,21 @@ Mock queryExecutor It.IsAny>())) .Returns(async () => { - // 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; + 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). @@ -1113,6 +1124,11 @@ await queryExecutor.Object.ExecuteQueryAsync( // 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."); } /// @@ -1149,7 +1165,7 @@ public async Task TestCancellationExceptionIsNotRetriedByRetryPolicy() Mock httpContextAccessor = new(); DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider); Mock queryExecutor - = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null); + = new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null, null); queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); @@ -1195,83 +1211,6 @@ await queryExecutor.Object.ExecuteQueryAsync( Assert.AreEqual(0, queryExecutorLogger.Invocations.Count); } - /// - /// Validates that when a request is cancelled during query execution, - /// the database execution time is still recorded in the HttpContext. - /// This ensures the finally block in ExecuteQueryAgainstDbAsync runs - /// even when the CancellationToken triggers an exception. - /// - [TestMethod, TestCategory(TestCategory.MSSQL)] - public async Task TestDbExecutionTimeRecordedEvenWhenRequestCancelled() - { - 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); - - queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary()); - - // Call the actual ExecuteQueryAgainstDbAsync to exercise the finally block. - queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync( - It.IsAny(), - It.IsAny(), - It.IsAny>(), - It.IsAny, Task>>(), - It.IsAny(), - It.IsAny(), - It.IsAny>())).CallBase(); - - // Create an HttpContext with a pre-cancelled CancellationToken. - CancellationTokenSource cts = new(); - cts.Cancel(); - context.RequestAborted = cts.Token; - - try - { - await queryExecutor.Object.ExecuteQueryAgainstDbAsync( - conn: null, - sqltext: string.Empty, - parameters: new Dictionary(), - dataReaderHandler: null, - dataSourceName: String.Empty, - httpContext: context, - args: null); - } - catch (Exception) - { - // SqlConnection is sealed and can't be mocked; conn is null so OpenAsync - // throws NullReferenceException. Ignore to verify the finally block behavior. - } - - // The finally block should have recorded execution time even though an exception occurred. - Assert.IsTrue( - context.Items.ContainsKey(TOTAL_DB_EXECUTION_TIME), - "HttpContext must contain the total db execution time even when the request is cancelled."); - } - [TestCleanup] public void CleanupAfterEachTest() { From d1b49fc4c26c386e207e5aca0d738bce899defb2 Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Tue, 24 Mar 2026 17:30:06 -0700 Subject: [PATCH 4/5] address PR comments --- src/Core/Resolvers/QueryExecutor.cs | 2 +- .../UnitTests/SqlQueryExecutorUnitTests.cs | 16 ++++++++++++---- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/Core/Resolvers/QueryExecutor.cs b/src/Core/Resolvers/QueryExecutor.cs index 2c1d080bef..d96b19af38 100644 --- a/src/Core/Resolvers/QueryExecutor.cs +++ b/src/Core/Resolvers/QueryExecutor.cs @@ -293,7 +293,7 @@ public virtual TConnection CreateConnection(string dataSourceName) TResult? result = default(TResult); try { - var commandBehavior = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ? CommandBehavior.SequentialAccess : 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); diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index 9ec1e536e3..30d0c99f84 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1120,10 +1120,18 @@ await queryExecutor.Object.ExecuteQueryAsync( 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 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( From 7f770be262ee459691e7e77b90af57e24b0fc9cd Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Thu, 26 Mar 2026 09:09:14 -0700 Subject: [PATCH 5/5] remove blank line --- src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index 807b6779cd..fa2e0e33f6 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1120,7 +1120,6 @@ await queryExecutor.Object.ExecuteQueryAsync( 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(