From 3ec5b55b778d09f383a848532acd81961042e4ad Mon Sep 17 00:00:00 2001 From: Nan Xing Date: Thu, 19 Mar 2026 05:16:19 -0700 Subject: [PATCH 1/3] =?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/3] 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/3] 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() {