diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 371391386e..d8a885a07b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -418,6 +418,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx inputParameters.ExecutionEnvironment, inputParameters.ReturnResultsInDeterministicOrder, inputParameters.ForcePassthrough, + inputParameters.EnableOptimisticDirectExecution, inputParameters.TestInjections); } @@ -752,7 +753,7 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (inputParameters.TestInjections == null || !inputParameters.TestInjections.EnableOptimisticDirectExecution) return null; + if (!inputParameters.EnableOptimisticDirectExecution) return null; // case 1: Is query going to a single partition bool hasPartitionKey = inputParameters.PartitionKey.HasValue @@ -824,6 +825,7 @@ public InputParameters( ExecutionEnvironment? executionEnvironment, bool? returnResultsInDeterministicOrder, bool forcePassthrough, + bool enableOptimisticDirectExecution, TestInjections testInjections) { this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); @@ -857,6 +859,7 @@ public InputParameters( this.ExecutionEnvironment = executionEnvironment.GetValueOrDefault(InputParameters.DefaultExecutionEnvironment); this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); this.ForcePassthrough = forcePassthrough; + this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution; this.TestInjections = testInjections; } @@ -873,6 +876,7 @@ public InputParameters( public bool ReturnResultsInDeterministicOrder { get; } public TestInjections TestInjections { get; } public bool ForcePassthrough { get; } + public bool EnableOptimisticDirectExecution { get; } public InputParameters WithContinuationToken(CosmosElement token) { @@ -889,6 +893,7 @@ public InputParameters WithContinuationToken(CosmosElement token) this.ExecutionEnvironment, this.ReturnResultsInDeterministicOrder, this.ForcePassthrough, + this.EnableOptimisticDirectExecution, this.TestInjections); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/TestInjections.cs b/Microsoft.Azure.Cosmos/src/Query/Core/TestInjections.cs index 7ee877bd14..1328d4dbab 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/TestInjections.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/TestInjections.cs @@ -13,16 +13,13 @@ public enum PipelineType OptimisticDirectExecution, } - public TestInjections(bool simulate429s, bool simulateEmptyPages, bool enableOptimisticDirectExecution = false, ResponseStats responseStats = null) + public TestInjections(bool simulate429s, bool simulateEmptyPages, ResponseStats responseStats = null) { this.SimulateThrottles = simulate429s; this.SimulateEmptyPages = simulateEmptyPages; this.Stats = responseStats; - this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution; } - public bool EnableOptimisticDirectExecution { get; } - public bool SimulateThrottles { get; } public bool SimulateEmptyPages { get; } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index 01a4e1b12e..a820a0a49a 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -144,6 +144,7 @@ public static QueryIterator Create( executionEnvironment: queryRequestOptions.ExecutionEnvironment, returnResultsInDeterministicOrder: queryRequestOptions.ReturnResultsInDeterministicOrder, forcePassthrough: forcePassthrough, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, testInjections: queryRequestOptions.TestSettings); return new QueryIterator( diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index 9e478c96b9..8711836191 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -43,6 +43,14 @@ public class QueryRequestOptions : RequestOptions /// public bool? EnableLowPrecisionOrderBy { get; set; } + /// + /// Gets or sets the option for customers to opt in for direct (optimistic) execution of the query.. + /// + /// + /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query. + /// + internal bool EnableOptimisticDirectExecution { get; set; } + /// /// Gets or sets the maximum number of items that can be buffered client side during /// parallel query execution in the Azure Cosmos DB service. diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs index acdb4d92de..9d89f6362a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs @@ -67,7 +67,8 @@ private static async Task TestPositiveOptimisticDirectExecutionOutput( QueryRequestOptions feedOptions = new QueryRequestOptions { MaxItemCount = -1, - TestSettings = GetTestInjections(simulate429s:false, simulateEmptyPages:false, enableOptimisticDirectExecution:true) + EnableOptimisticDirectExecution = true, + TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: false, new TestInjections.ResponseStats()) }; // check if pipeline returns empty continuation token @@ -172,13 +173,11 @@ private static async Task TestPositiveOptimisticDirectExecutionOutput( feedOptions = new QueryRequestOptions { MaxItemCount = pageSizeOptions[i], - PartitionKey = queryAndResults[j].PartitionKey == null + PartitionKey = queryAndResults[j].PartitionKey == null ? null : new Cosmos.PartitionKey(queryAndResults[j].PartitionKey), - TestSettings = GetTestInjections( - simulate429s: false, - simulateEmptyPages: false, - enableOptimisticDirectExecution: queryAndResults[j].EnableOptimisticDirectExecution) + EnableOptimisticDirectExecution = queryAndResults[j].EnableOptimisticDirectExecution, + TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: false, new TestInjections.ResponseStats()) }; List items = await RunQueryAsync( @@ -208,7 +207,8 @@ private static async Task TestNegativeOptimisticDirectExecutionOutput( QueryRequestOptions feedOptions = new QueryRequestOptions { PartitionKey = new Cosmos.PartitionKey("/value"), - TestSettings = GetTestInjections(simulate429s: false, simulateEmptyPages: false, enableOptimisticDirectExecution: true) + EnableOptimisticDirectExecution = true, + TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: false, new TestInjections.ResponseStats()) }; // check if bad continuation queries and syntax error queries are handled by pipeline @@ -241,15 +241,6 @@ await container.GetItemQueryIterator( } } - private static TestInjections GetTestInjections(bool simulate429s, bool simulateEmptyPages, bool enableOptimisticDirectExecution) - { - return new TestInjections( - simulate429s, - simulateEmptyPages, - enableOptimisticDirectExecution, - new TestInjections.ResponseStats()); - } - private struct SinglePartitionWithContinuationsArgs { public int NumberOfDocuments; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs index 4775716efa..c196af9d44 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs @@ -597,6 +597,7 @@ internal static async Task> QueryWithCosmosElementContinuationTokenAsync Properties = queryRequestOptions.Properties, IsEffectivePartitionKeyRouting = queryRequestOptions.IsEffectivePartitionKeyRouting, CosmosElementContinuationToken = queryRequestOptions.CosmosElementContinuationToken, + EnableOptimisticDirectExecution = queryRequestOptions.EnableOptimisticDirectExecution, TestSettings = queryRequestOptions.TestSettings, }; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 6737a3d338..b6a78c111d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -148,11 +148,12 @@ public void NegativeOptimisticDirectExecutionOutput() // This test confirms that TestInjection.EnableOptimisticDirectExection is set to false from default. // Check test "TestPipelineForDistributedQueryAsync" to understand why this is done [TestMethod] - public async Task TestDefaultTestInjectionSettingsAsync() + public async Task TestDefaultQueryRequestOptionsSettings() { - TestInjections testInjection = new TestInjections(simulate429s: false, simulateEmptyPages: false); - Assert.AreEqual(testInjection.EnableOptimisticDirectExecution, false); + QueryRequestOptions requestOptions = new QueryRequestOptions(); + + Assert.AreEqual(requestOptions.EnableOptimisticDirectExecution, false); } // test checks that the pipeline can take a query to the backend and returns its associated document(s). @@ -563,6 +564,7 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect executionEnvironment: null, returnResultsInDeterministicOrder: null, forcePassthrough: false, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, testInjections: queryRequestOptions.TestSettings); string databaseId = "db1234"; @@ -587,7 +589,8 @@ private static QueryRequestOptions GetQueryRequestOptions(bool enableOptimisticD { MaxConcurrency = 0, MaxItemCount = 10, - TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, enableOptimisticDirectExecution: enableOptimisticDirectExecution, new TestInjections.ResponseStats()), + EnableOptimisticDirectExecution = enableOptimisticDirectExecution, + TestSettings = new TestInjections(simulate429s: true, simulateEmptyPages: false, new TestInjections.ResponseStats()), Properties = new Dictionary() { { HttpConstants.HttpHeaders.EnumerationDirection, ""},