Skip to content

Commit 290eff4

Browse files
naxing123Aniruddh25anushakolan
authored
Cherry-pick #3300: Use cancellation token in DbCommand.ExecuteReaderAsync call in DAB layer (#3406)
## Why make this change? This backport fix for to address GraphQL timeout related issues. ## What is this change? Cherry-picked PR: Use cancellation token in DbCommand.ExecuteReaderAsync call in DAB layer #3300 ## How was this tested? 1. Existing and newly added tests in the source PR. 2. 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. ## Sample Request(s) N/A --------- Co-authored-by: Aniruddh Munde <[email protected]> Co-authored-by: Anusha Kolan <[email protected]>
1 parent a9d8494 commit 290eff4

2 files changed

Lines changed: 209 additions & 2 deletions

File tree

‎src/Core/Resolvers/QueryExecutor.cs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -293,8 +293,11 @@ public virtual TConnection CreateConnection(string dataSourceName)
293293
TResult? result = default(TResult);
294294
try
295295
{
296-
using DbDataReader dbDataReader = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ?
297-
await cmd.ExecuteReaderAsync(CommandBehavior.SequentialAccess) : await cmd.ExecuteReaderAsync(CommandBehavior.CloseConnection);
296+
CommandBehavior commandBehavior = ConfigProvider.GetConfig().MaxResponseSizeLogicEnabled() ? CommandBehavior.SequentialAccess : CommandBehavior.CloseConnection;
297+
// 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.
298+
CancellationToken cancellationToken = httpContext?.RequestAborted ?? CancellationToken.None;
299+
using DbDataReader dbDataReader = await cmd.ExecuteReaderAsync(commandBehavior, cancellationToken);
300+
298301
if (dataReaderHandler is not null && dbDataReader is not null)
299302
{
300303
result = await dataReaderHandler(dbDataReader, args);

‎src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
using System.Linq;
1212
using System.Net;
1313
using System.Reflection;
14+
using System.Threading;
1415
using System.Threading.Tasks;
1516
using Azure.Core;
1617
using Azure.DataApiBuilder.Config;
@@ -669,6 +670,209 @@ public void ValidateStreamingLogicForEmptyCellsAsync()
669670
Assert.AreEqual(availableSize, (int)runtimeConfig.MaxResponseSizeMB() * 1024 * 1024);
670671
}
671672

673+
/// <summary>
674+
/// Validates that when the CancellationToken from httpContext.RequestAborted times out
675+
/// during a long-running query execution (simulating ExecuteReaderAsync being interrupted
676+
/// by a token timeout), the resulting TaskCanceledException propagates through the Polly
677+
/// retry policy without any retry attempts.
678+
/// Unlike TestCancellationExceptionIsNotRetriedByRetryPolicy which throws immediately,
679+
/// this test simulates a real timeout where the cancellation occurs asynchronously
680+
/// after a delay.
681+
/// </summary>
682+
[TestMethod, TestCategory(TestCategory.MSSQL)]
683+
public async Task TestCancellationTokenTimeoutDuringQueryExecutionAsync()
684+
{
685+
RuntimeConfig mockConfig = new(
686+
Schema: "",
687+
DataSource: new(DatabaseType.MSSQL, "", new()),
688+
Runtime: new(
689+
Rest: new(),
690+
GraphQL: new(),
691+
Mcp: new(),
692+
Host: new(null, null)
693+
),
694+
Entities: new(new Dictionary<string, Entity>())
695+
);
696+
697+
MockFileSystem fileSystem = new();
698+
fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(mockConfig.ToJson()));
699+
FileSystemRuntimeConfigLoader loader = new(fileSystem);
700+
RuntimeConfigProvider provider = new(loader)
701+
{
702+
IsLateConfigured = true
703+
};
704+
705+
Mock<ILogger<QueryExecutor<SqlConnection>>> queryExecutorLogger = new();
706+
Mock<IHttpContextAccessor> httpContextAccessor = new();
707+
HttpContext context = new DefaultHttpContext();
708+
httpContextAccessor.Setup(x => x.HttpContext).Returns(context);
709+
DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider);
710+
Mock<MsSqlQueryExecutor> queryExecutor
711+
= new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null);
712+
713+
queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary<string, DbConnectionStringBuilder>());
714+
715+
queryExecutor.Setup(x => x.CreateConnection(
716+
It.IsAny<string>())).CallBase();
717+
718+
// Set up a CancellationTokenSource that times out after a short delay,
719+
// simulating httpContext.RequestAborted firing due to a client timeout.
720+
CancellationTokenSource cts = new();
721+
cts.CancelAfter(TimeSpan.FromMilliseconds(100));
722+
context.RequestAborted = cts.Token;
723+
724+
// Mock ExecuteQueryAgainstDbAsync to simulate a long-running database query
725+
// that is interrupted when the CancellationToken times out.
726+
// Task.Delay with the cancellation token throws TaskCanceledException when the
727+
// token fires, mimicking cmd.ExecuteReaderAsync being cancelled by a timed-out token.
728+
// The Stopwatch + finally block mirrors the real ExecuteQueryAgainstDbAsync to verify
729+
// that execution time is recorded even when a timeout occurs.
730+
queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync(
731+
It.IsAny<SqlConnection>(),
732+
It.IsAny<string>(),
733+
It.IsAny<IDictionary<string, DbConnectionParam>>(),
734+
It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
735+
It.IsAny<HttpContext>(),
736+
provider.GetConfig().DefaultDataSourceName,
737+
It.IsAny<List<string>>()))
738+
.Returns(async () =>
739+
{
740+
Stopwatch timer = Stopwatch.StartNew();
741+
try
742+
{
743+
// Simulate a long-running query interrupted by token timeout.
744+
// Timeout.Infinite (-1) means "wait forever" — the only way this
745+
// completes is when cts.Token fires after ~100 ms, which causes
746+
// Task.Delay to throw TaskCanceledException.
747+
await Task.Delay(Timeout.Infinite, cts.Token);
748+
return (object)null;
749+
}
750+
finally
751+
{
752+
timer.Stop();
753+
queryExecutor.Object.AddDbExecutionTimeToMiddlewareContext(timer.ElapsedMilliseconds);
754+
}
755+
});
756+
757+
// Call the actual ExecuteQueryAsync method (includes Polly retry policy).
758+
queryExecutor.Setup(x => x.ExecuteQueryAsync(
759+
It.IsAny<string>(),
760+
It.IsAny<IDictionary<string, DbConnectionParam>>(),
761+
It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
762+
It.IsAny<string>(),
763+
It.IsAny<HttpContext>(),
764+
It.IsAny<List<string>>())).CallBase();
765+
766+
// Act & Assert: TaskCanceledException should propagate without retries.
767+
await Assert.ThrowsExceptionAsync<TaskCanceledException>(async () =>
768+
{
769+
await queryExecutor.Object.ExecuteQueryAsync<object>(
770+
sqltext: string.Empty,
771+
parameters: new Dictionary<string, DbConnectionParam>(),
772+
dataReaderHandler: null,
773+
dataSourceName: String.Empty,
774+
httpContext: context,
775+
args: null);
776+
});
777+
778+
// Verify that the underlying database execution is invoked exactly once,
779+
// confirming that Polly does not perform any retries for TaskCanceledException.
780+
queryExecutor.Verify(q => q.ExecuteQueryAgainstDbAsync(
781+
It.IsAny<SqlConnection>(),
782+
It.IsAny<string>(),
783+
It.IsAny<IDictionary<string, DbConnectionParam>>(),
784+
It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
785+
It.IsAny<HttpContext>(),
786+
provider.GetConfig().DefaultDataSourceName,
787+
It.IsAny<List<string>>()),
788+
Times.Once);
789+
790+
// Verify the finally block recorded execution time even though the token timed out.
791+
Assert.IsTrue(
792+
context.Items.ContainsKey(TOTAL_DB_EXECUTION_TIME),
793+
"HttpContext must contain the total db execution time even when the request is cancelled.");
794+
}
795+
796+
/// <summary>
797+
/// Validates that when ExecuteQueryAgainstDbAsync throws OperationCanceledException
798+
/// (e.g., due to client disconnect via httpContext.RequestAborted cancellation token),
799+
/// the Polly retry policy does NOT retry and the exception propagates to the caller.
800+
/// The retry policy is configured to only handle DbException, so OperationCanceledException
801+
/// should be immediately re-thrown without any retry attempts.
802+
/// </summary>
803+
[TestMethod, TestCategory(TestCategory.MSSQL)]
804+
public async Task TestCancellationExceptionIsNotRetriedByRetryPolicy()
805+
{
806+
RuntimeConfig mockConfig = new(
807+
Schema: "",
808+
DataSource: new(DatabaseType.MSSQL, "", new()),
809+
Runtime: new(
810+
Rest: new(),
811+
GraphQL: new(),
812+
Mcp: new(),
813+
Host: new(null, null)
814+
),
815+
Entities: new(new Dictionary<string, Entity>())
816+
);
817+
818+
MockFileSystem fileSystem = new();
819+
fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(mockConfig.ToJson()));
820+
FileSystemRuntimeConfigLoader loader = new(fileSystem);
821+
RuntimeConfigProvider provider = new(loader)
822+
{
823+
IsLateConfigured = true
824+
};
825+
826+
Mock<ILogger<QueryExecutor<SqlConnection>>> queryExecutorLogger = new();
827+
Mock<IHttpContextAccessor> httpContextAccessor = new();
828+
DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(provider);
829+
Mock<MsSqlQueryExecutor> queryExecutor
830+
= new(provider, dbExceptionParser, queryExecutorLogger.Object, httpContextAccessor.Object, null);
831+
832+
queryExecutor.Setup(x => x.ConnectionStringBuilders).Returns(new Dictionary<string, DbConnectionStringBuilder>());
833+
834+
queryExecutor.Setup(x => x.CreateConnection(
835+
It.IsAny<string>())).CallBase();
836+
837+
// Mock ExecuteQueryAgainstDbAsync to throw OperationCanceledException,
838+
// simulating a cancelled CancellationToken from httpContext.RequestAborted.
839+
queryExecutor.Setup(x => x.ExecuteQueryAgainstDbAsync(
840+
It.IsAny<SqlConnection>(),
841+
It.IsAny<string>(),
842+
It.IsAny<IDictionary<string, DbConnectionParam>>(),
843+
It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
844+
It.IsAny<HttpContext>(),
845+
provider.GetConfig().DefaultDataSourceName,
846+
It.IsAny<List<string>>()))
847+
.ThrowsAsync(new OperationCanceledException("The operation was canceled."));
848+
849+
// Call the actual ExecuteQueryAsync method.
850+
queryExecutor.Setup(x => x.ExecuteQueryAsync(
851+
It.IsAny<string>(),
852+
It.IsAny<IDictionary<string, DbConnectionParam>>(),
853+
It.IsAny<Func<DbDataReader, List<string>, Task<object>>>(),
854+
It.IsAny<string>(),
855+
It.IsAny<HttpContext>(),
856+
It.IsAny<List<string>>())).CallBase();
857+
858+
// Act & Assert: OperationCanceledException should propagate without retries.
859+
await Assert.ThrowsExceptionAsync<OperationCanceledException>(async () =>
860+
{
861+
await queryExecutor.Object.ExecuteQueryAsync<object>(
862+
sqltext: string.Empty,
863+
parameters: new Dictionary<string, DbConnectionParam>(),
864+
dataReaderHandler: null,
865+
dataSourceName: String.Empty,
866+
httpContext: null,
867+
args: null);
868+
});
869+
870+
// Verify no retry log messages were emitted. Since IsLateConfigured is true,
871+
// the debug log is skipped, and since Polly doesn't handle OperationCanceledException,
872+
// no retry occurs → zero logger invocations.
873+
Assert.AreEqual(0, queryExecutorLogger.Invocations.Count);
874+
}
875+
672876
[TestCleanup]
673877
public void CleanupAfterEachTest()
674878
{

0 commit comments

Comments
 (0)