From c0c912df468a69ceb621d5963de024273e353d6e Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 25 Jul 2023 13:57:29 -0700 Subject: [PATCH 01/11] Added code to throw exception if ODE continuation token goes into non ODE pipeline --- .../CosmosQueryExecutionContextFactory.cs | 14 +++++ .../CosmosBasicQueryTests.cs | 63 ++++++++++++++++++- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 47aa1c4490..e51233943f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -32,6 +32,7 @@ internal static class CosmosQueryExecutionContextFactory { private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition"; private const string OptimisticDirectExecution = "OptimisticDirectExecution"; + private const string OptimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; private const string Passthrough = "Passthrough"; private const string Specialized = "Specialized"; private const int PageSizeFactorForTop = 5; @@ -307,6 +308,19 @@ private static async Task> TryCreateFromPartitione } else { + if (inputParameters.InitialUserContinuationToken != null) + { + CosmosObject cosmosObjectContinuationToken = inputParameters.InitialUserContinuationToken as CosmosObject; + if (cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Operation cannot be resumed with the following token, as it requires the Optimistic Direct Execution pipeline which has been disabled on your SDK. " + + $"Enable Optimistic Direct Execution from QueryRequestOptions to avoid this issue." + + $"{inputParameters.InitialUserContinuationToken}.")); + } + } + if (createPassthroughQuery) { SetTestInjectionPipelineType(inputParameters, Passthrough); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index 25ac9a705f..cbb16f127f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -15,7 +15,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Cosmos.Scripts; using Microsoft.Azure.Cosmos.Fluent; using Microsoft.Azure.Cosmos.Linq; - using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents.Collections; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -786,6 +785,68 @@ public async Task QueryActivityIdWithContinuationTokenAndTraceTest() } + [TestMethod] + public async Task TesOdeTokenCompatibilityWithSpecializedPipeline() + { + using (ITrace rootTrace = Trace.GetRootTrace("Root Trace")) + { + ResponseMessage responseMessage = null; + int count = 0; + string query = "select top 200 * from c"; + string expectedErrorMessage = "Operation cannot be resumed with the following token, as it requires Optimistic Direct Execution pipeline, which has been disabled for your SDK"; + + CosmosClient client = DirectCosmosClient; + Container container = client.GetContainer(DatabaseId, ContainerId); + // Create items + for (int i = 0; i < 500; i++) + { + await container.CreateItemAsync(ToDoActivity.CreateRandomToDoActivity()); + } + + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxItemCount = 50, + EnableOptimisticDirectExecution = true + }; + + FeedIteratorInternal feedIterator = + (FeedIteratorInternal)container.GetItemQueryStreamIterator( + query, + null, + queryRequestOptions); + + responseMessage = await feedIterator.ReadNextAsync(rootTrace, CancellationToken.None); + string continuationToken = responseMessage.ContinuationToken; + Collection response = TestCommon.SerializerCore.FromStream>(responseMessage.Content).Data; + count += response.Count; + + QueryRequestOptions newQueryRequestOptions = new QueryRequestOptions + { + MaxItemCount = 50, + EnableOptimisticDirectExecution = false + }; + + // use Continuation Token to create new iterator and use same trace + FeedIteratorInternal feedIteratorNew = + (FeedIteratorInternal)container.GetItemQueryStreamIterator( + query, + continuationToken, + newQueryRequestOptions); + + while (feedIteratorNew.HasMoreResults) + { + responseMessage = await feedIteratorNew.ReadNextAsync(rootTrace, CancellationToken.None); + if (responseMessage.Content != null) + { + Collection newResponse = TestCommon.SerializerCore.FromStream>(responseMessage.Content).Data; + count += newResponse.Count; + } + } + + Assert.IsTrue(responseMessage.CosmosException.ToString().Contains(expectedErrorMessage)); + } + } + private class CustomHandler : RequestHandler { string correlatedActivityId; From 5f4453f4c5489e8d90fcedd6076c342131f43a65 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 25 Jul 2023 14:19:08 -0700 Subject: [PATCH 02/11] Removed count variable --- .../CosmosBasicQueryTests.cs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index cbb16f127f..9d1375961d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -791,12 +791,12 @@ public async Task TesOdeTokenCompatibilityWithSpecializedPipeline() using (ITrace rootTrace = Trace.GetRootTrace("Root Trace")) { ResponseMessage responseMessage = null; - int count = 0; string query = "select top 200 * from c"; string expectedErrorMessage = "Operation cannot be resumed with the following token, as it requires Optimistic Direct Execution pipeline, which has been disabled for your SDK"; CosmosClient client = DirectCosmosClient; Container container = client.GetContainer(DatabaseId, ContainerId); + // Create items for (int i = 0; i < 500; i++) { @@ -817,8 +817,6 @@ public async Task TesOdeTokenCompatibilityWithSpecializedPipeline() responseMessage = await feedIterator.ReadNextAsync(rootTrace, CancellationToken.None); string continuationToken = responseMessage.ContinuationToken; - Collection response = TestCommon.SerializerCore.FromStream>(responseMessage.Content).Data; - count += response.Count; QueryRequestOptions newQueryRequestOptions = new QueryRequestOptions { @@ -836,11 +834,6 @@ public async Task TesOdeTokenCompatibilityWithSpecializedPipeline() while (feedIteratorNew.HasMoreResults) { responseMessage = await feedIteratorNew.ReadNextAsync(rootTrace, CancellationToken.None); - if (responseMessage.Content != null) - { - Collection newResponse = TestCommon.SerializerCore.FromStream>(responseMessage.Content).Data; - count += newResponse.Count; - } } Assert.IsTrue(responseMessage.CosmosException.ToString().Contains(expectedErrorMessage)); From dfb902c810a22974472fc5b94ce73faba21e1c6a Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Tue, 25 Jul 2023 15:03:55 -0700 Subject: [PATCH 03/11] Updated test name --- .../CosmosBasicQueryTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index 9d1375961d..d2756e71fb 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -786,7 +786,7 @@ public async Task QueryActivityIdWithContinuationTokenAndTraceTest() } [TestMethod] - public async Task TesOdeTokenCompatibilityWithSpecializedPipeline() + public async Task TesOdeTokenCompatibilityWithNonOdePipeline() { using (ITrace rootTrace = Trace.GetRootTrace("Root Trace")) { From 835f8bde265223e8a67f1d5e77779f58b87818cb Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 26 Jul 2023 14:34:49 -0700 Subject: [PATCH 04/11] Removed ODE continuation token logic from caller class --- .../CosmosQueryExecutionContextFactory.cs | 26 +++++++++---------- ...imisticDirectExecutionContinuationToken.cs | 23 +++++++++++----- .../CosmosBasicQueryTests.cs | 2 +- 3 files changed, 30 insertions(+), 21 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index e51233943f..c505a8e753 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -308,19 +308,6 @@ private static async Task> TryCreateFromPartitione } else { - if (inputParameters.InitialUserContinuationToken != null) - { - CosmosObject cosmosObjectContinuationToken = inputParameters.InitialUserContinuationToken as CosmosObject; - if (cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Operation cannot be resumed with the following token, as it requires the Optimistic Direct Execution pipeline which has been disabled on your SDK. " + - $"Enable Optimistic Direct Execution from QueryRequestOptions to avoid this issue." + - $"{inputParameters.InitialUserContinuationToken}.")); - } - } - if (createPassthroughQuery) { SetTestInjectionPipelineType(inputParameters, Passthrough); @@ -766,7 +753,18 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (!inputParameters.EnableOptimisticDirectExecution) return null; + if (!inputParameters.EnableOptimisticDirectExecution) + { + if (inputParameters.InitialUserContinuationToken != null) + { + if (OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) + { + throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " + + $"{inputParameters.InitialUserContinuationToken}"); + } + } + + } Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!"); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index 59f87295f3..37d6812350 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -4,9 +4,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery { - using System; using System.Collections.Generic; - using Microsoft.Azure.Cosmos.ChangeFeed; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.Query.Core.Exceptions; using Microsoft.Azure.Cosmos.Query.Core.Monads; @@ -30,6 +28,19 @@ public OptimisticDirectExecutionContinuationToken(ParallelContinuationToken toke public Range Range => this.Token.Range; + public static bool IsOptimisticDirectExecutionContinuationToken(CosmosElement continuationToken) + { + CosmosObject cosmosObjectContinuationToken = continuationToken as CosmosObject; + if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) + { + return false; + } + else + { + return true; + } + } + public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuationToken continuationToken) { CosmosElement inner = ParallelContinuationToken.ToCosmosElement(continuationToken.Token); @@ -42,14 +53,14 @@ public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuatio public static TryCatch TryCreateFromCosmosElement(CosmosElement cosmosElement) { - CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject; - if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) + if (!IsOptimisticDirectExecutionContinuationToken(cosmosElement)) { return TryCatch.FromException( - new MalformedContinuationTokenException( - message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n")); + new MalformedContinuationTokenException( + message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n")); } + CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject; TryCatch inner = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosObjectContinuationToken[OptimisticDirectExecutionToken]); return inner.Succeeded ? diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index d2756e71fb..c2de3b5a7a 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -792,7 +792,7 @@ public async Task TesOdeTokenCompatibilityWithNonOdePipeline() { ResponseMessage responseMessage = null; string query = "select top 200 * from c"; - string expectedErrorMessage = "Operation cannot be resumed with the following token, as it requires Optimistic Direct Execution pipeline, which has been disabled for your SDK"; + string expectedErrorMessage = "The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. "; CosmosClient client = DirectCosmosClient; Container container = client.GetContainer(DatabaseId, ContainerId); From 4e8cb4ff04c5e017aa3c9239cb7035a7e5f884b3 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Fri, 28 Jul 2023 10:51:37 -0700 Subject: [PATCH 05/11] Simplified code --- .../CosmosQueryExecutionContextFactory.cs | 15 +++++---------- .../OptimisticDirectExecutionContinuationToken.cs | 9 +-------- 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index c505a8e753..88822b23ff 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -753,17 +753,12 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (!inputParameters.EnableOptimisticDirectExecution) + if (!inputParameters.EnableOptimisticDirectExecution + && inputParameters.InitialUserContinuationToken != null + && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) { - if (inputParameters.InitialUserContinuationToken != null) - { - if (OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) - { - throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " + - $"{inputParameters.InitialUserContinuationToken}"); - } - } - + throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " + + $"{inputParameters.InitialUserContinuationToken}"); } Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!"); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index 37d6812350..c89d4aaee3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -31,14 +31,7 @@ public OptimisticDirectExecutionContinuationToken(ParallelContinuationToken toke public static bool IsOptimisticDirectExecutionContinuationToken(CosmosElement continuationToken) { CosmosObject cosmosObjectContinuationToken = continuationToken as CosmosObject; - if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)) - { - return false; - } - else - { - return true; - } + return !(cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)); } public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuationToken continuationToken) From 9be3def61839078a62e62ed9943d01714729c7d0 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 2 Aug 2023 11:08:41 -0700 Subject: [PATCH 06/11] Fixed comments --- .../CosmosQueryExecutionContextFactory.cs | 12 ++- ...imisticDirectExecutionContinuationToken.cs | 4 +- .../CosmosBasicQueryTests.cs | 79 +++++++++---------- 3 files changed, 47 insertions(+), 48 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 88822b23ff..bc5d25185f 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -753,12 +753,16 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP ContainerQueryProperties containerQueryProperties, ITrace trace) { - if (!inputParameters.EnableOptimisticDirectExecution - && inputParameters.InitialUserContinuationToken != null - && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) + if (!inputParameters.EnableOptimisticDirectExecution) { - throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " + + if (inputParameters.InitialUserContinuationToken != null + && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) + { + throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " + $"{inputParameters.InitialUserContinuationToken}"); + } + + return null; } Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!"); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index c89d4aaee3..832ff1031b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -30,7 +30,7 @@ public OptimisticDirectExecutionContinuationToken(ParallelContinuationToken toke public static bool IsOptimisticDirectExecutionContinuationToken(CosmosElement continuationToken) { - CosmosObject cosmosObjectContinuationToken = continuationToken as CosmosObject; + CosmosObject cosmosObjectContinuationToken = (CosmosObject)continuationToken; return !(cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)); } @@ -53,7 +53,7 @@ public static TryCatch TryCreateFrom message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n")); } - CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject; + CosmosObject cosmosObjectContinuationToken = (CosmosObject)cosmosElement; TryCatch inner = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosObjectContinuationToken[OptimisticDirectExecutionToken]); return inner.Succeeded ? diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index c2de3b5a7a..5dd8cc61c8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -788,56 +788,51 @@ public async Task QueryActivityIdWithContinuationTokenAndTraceTest() [TestMethod] public async Task TesOdeTokenCompatibilityWithNonOdePipeline() { - using (ITrace rootTrace = Trace.GetRootTrace("Root Trace")) - { - ResponseMessage responseMessage = null; - string query = "select top 200 * from c"; - string expectedErrorMessage = "The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. "; - - CosmosClient client = DirectCosmosClient; - Container container = client.GetContainer(DatabaseId, ContainerId); + string query = "select top 200 * from c"; + CosmosClient client = DirectCosmosClient; + Container container = client.GetContainer(DatabaseId, ContainerId); - // Create items - for (int i = 0; i < 500; i++) - { - await container.CreateItemAsync(ToDoActivity.CreateRandomToDoActivity()); - } - - QueryRequestOptions queryRequestOptions = new QueryRequestOptions - { - MaxItemCount = 50, - EnableOptimisticDirectExecution = true - }; + // Create items + for (int i = 0; i < 500; i++) + { + await container.CreateItemAsync(ToDoActivity.CreateRandomToDoActivity()); + } - FeedIteratorInternal feedIterator = - (FeedIteratorInternal)container.GetItemQueryStreamIterator( - query, - null, - queryRequestOptions); + QueryRequestOptions queryRequestOptions = new QueryRequestOptions + { + MaxItemCount = 50, + EnableOptimisticDirectExecution = true + }; - responseMessage = await feedIterator.ReadNextAsync(rootTrace, CancellationToken.None); - string continuationToken = responseMessage.ContinuationToken; + FeedIteratorInternal feedIterator = + (FeedIteratorInternal)container.GetItemQueryStreamIterator( + query, + null, + queryRequestOptions); - QueryRequestOptions newQueryRequestOptions = new QueryRequestOptions - { - MaxItemCount = 50, - EnableOptimisticDirectExecution = false - }; + ResponseMessage responseMessage = await feedIterator.ReadNextAsync(CancellationToken.None); + string continuationToken = responseMessage.ContinuationToken; - // use Continuation Token to create new iterator and use same trace - FeedIteratorInternal feedIteratorNew = - (FeedIteratorInternal)container.GetItemQueryStreamIterator( - query, - continuationToken, - newQueryRequestOptions); + QueryRequestOptions newQueryRequestOptions = new QueryRequestOptions + { + MaxItemCount = 50, + EnableOptimisticDirectExecution = false + }; - while (feedIteratorNew.HasMoreResults) - { - responseMessage = await feedIteratorNew.ReadNextAsync(rootTrace, CancellationToken.None); - } + // use Continuation Token to create new iterator and use same trace + FeedIterator feedIteratorNew = + container.GetItemQueryStreamIterator( + query, + continuationToken, + newQueryRequestOptions); - Assert.IsTrue(responseMessage.CosmosException.ToString().Contains(expectedErrorMessage)); + while (feedIteratorNew.HasMoreResults) + { + responseMessage = await feedIteratorNew.ReadNextAsync(CancellationToken.None); } + + string expectedErrorMessage = "The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. "; + Assert.IsTrue(responseMessage.CosmosException.ToString().Contains(expectedErrorMessage)); } private class CustomHandler : RequestHandler From 603dbdd9854c67eb250217669beded8b71f574c8 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 2 Aug 2023 12:34:19 -0700 Subject: [PATCH 07/11] Updated continuation token cast --- .../OptimisticDirectExecutionContinuationToken.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs index 832ff1031b..76b2531f01 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs @@ -30,7 +30,7 @@ public OptimisticDirectExecutionContinuationToken(ParallelContinuationToken toke public static bool IsOptimisticDirectExecutionContinuationToken(CosmosElement continuationToken) { - CosmosObject cosmosObjectContinuationToken = (CosmosObject)continuationToken; + CosmosObject cosmosObjectContinuationToken = continuationToken as CosmosObject; return !(cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken)); } From 46e2b331cd47169cc3d924ac03d2d3c2cd143a99 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Wed, 2 Aug 2023 14:49:45 -0700 Subject: [PATCH 08/11] Removed const string for continuation token --- .../Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index bc5d25185f..43ebf6a2d8 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -32,7 +32,6 @@ internal static class CosmosQueryExecutionContextFactory { private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition"; private const string OptimisticDirectExecution = "OptimisticDirectExecution"; - private const string OptimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; private const string Passthrough = "Passthrough"; private const string Specialized = "Specialized"; private const int PageSizeFactorForTop = 5; From f9b6a97060cffa4cf045efb54a7e48c5a1057246 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 3 Aug 2023 10:37:19 -0700 Subject: [PATCH 09/11] Added Ignore flag for test --- .../CosmosBasicQueryTests.cs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs index 5dd8cc61c8..bfd6366104 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosBasicQueryTests.cs @@ -785,6 +785,8 @@ public async Task QueryActivityIdWithContinuationTokenAndTraceTest() } + //TODO: Remove Ignore flag once emulator is updated to 0415 + [Ignore] [TestMethod] public async Task TesOdeTokenCompatibilityWithNonOdePipeline() { From 5ef5b7fea8f21e3ff10f5f1ddd042d31429b39d6 Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 3 Aug 2023 14:51:47 -0700 Subject: [PATCH 10/11] Added baseline test --- ...misticDirectExecutionQueryBaselineTests.cs | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) 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 81c765adc4..26a1ea4a0f 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 @@ -19,6 +19,7 @@ using Microsoft.Azure.Cosmos.Query.Core.Pipeline; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; @@ -190,6 +191,52 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() Assert.AreEqual(100, documentCountInSinglePartition); } + [TestMethod] + public async Task TestOdeTokenWithSpecializedPipeline() + { + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: Guid.NewGuid().ToString(), + range: new Documents.Routing.Range("A", "B", true, false)); + + OptimisticDirectExecutionContinuationToken optimisticDirectExecutionContinuationToken = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); + CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(optimisticDirectExecutionContinuationToken); + + OptimisticDirectExecutionTestInput input = CreateInput( + description: @"Single Partition Key and Value Field", + query: "SELECT VALUE COUNT(1) FROM c", + expectedOptimisticDirectExecution: false, + partitionKeyPath: @"/pk", + partitionKeyValue: "a", + continuationToken: cosmosElementContinuationToken); + + IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(input.PartitionKeyDefinition); + DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); + + QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: input.ExpectedOptimisticDirectExecution); + + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); + + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + + string expectedErrorMessage = "The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. "; + + while (await queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton)) + { + if (queryPipelineStage.Current.Failed) + { + Assert.IsTrue(queryPipelineStage.Current.InnerMostException.ToString().Contains(expectedErrorMessage)); + return; + } + + Assert.IsFalse(true); + break; + } + } + [TestMethod] public async Task TestQueriesWhichNeverRequireDistribution() { From 98f99f03bc71ad0a88bf0a03da6541f711c66edd Mon Sep 17 00:00:00 2001 From: Aditya Kotalwar Date: Thu, 3 Aug 2023 14:56:58 -0700 Subject: [PATCH 11/11] Updated baseline test --- .../Query/OptimisticDirectExecutionQueryBaselineTests.cs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 26a1ea4a0f..5c0e66f1a6 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 @@ -194,6 +194,7 @@ public async Task TestPipelineForBackendDocumentsOnSinglePartitionAsync() [TestMethod] public async Task TestOdeTokenWithSpecializedPipeline() { + int numItems = 100; ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( token: Guid.NewGuid().ToString(), range: new Documents.Routing.Range("A", "B", true, false)); @@ -209,13 +210,9 @@ public async Task TestOdeTokenWithSpecializedPipeline() partitionKeyValue: "a", continuationToken: cosmosElementContinuationToken); - IMonadicDocumentContainer monadicDocumentContainer = new InMemoryContainer(input.PartitionKeyDefinition); - DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); - + DocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems, multiPartition: false); QueryRequestOptions queryRequestOptions = GetQueryRequestOptions(enableOptimisticDirectExecution: input.ExpectedOptimisticDirectExecution); - (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = CreateInputParamsAndQueryContext(input, queryRequestOptions); - IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( documentContainer, cosmosQueryContextCore,