From 4bbfd41b6b855a17108953cfe4a549dcc93caa11 Mon Sep 17 00:00:00 2001 From: Jason Thorsness Date: Mon, 21 Oct 2019 15:23:12 -0700 Subject: [PATCH 1/3] fix error with partial count returned --- ...smosCrossPartitionQueryExecutionContext.cs | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index b3bd055cf8..b955a3d19b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -218,7 +218,7 @@ protected abstract string ContinuationToken /// /// Gets a value indicating whether the context still has more results. /// - private bool HasMoreResults => this.itemProducerForest.Count != 0 && this.CurrentItemProducerTree().HasMoreResults; + private bool HasMoreResults => this.FailureResponse != null || (this.itemProducerForest.Count != 0 && this.CurrentItemProducerTree().HasMoreResults); /// /// Gets the number of documents we can still buffer. @@ -353,7 +353,15 @@ public override async Task DrainAsync(int maxElements, Cancel if (this.FailureResponse != null) { this.Stop(); - return this.FailureResponse.Value; + + try + { + return this.FailureResponse; + } + finally + { + this.FailureResponse = null; + } } // Drain the results. If there is no results and a failure then return the failure. @@ -361,7 +369,15 @@ public override async Task DrainAsync(int maxElements, Cancel if ((results == null || results.Count == 0) && this.FailureResponse != null) { this.Stop(); - return this.FailureResponse.Value; + + try + { + return this.FailureResponse; + } + finally + { + this.FailureResponse = null; + } } string continuation = this.ContinuationToken; From 0ec79d754d30c54fae21422057e6d712da7721f4 Mon Sep 17 00:00:00 2001 From: Jake Willey Date: Tue, 22 Oct 2019 06:12:15 -0700 Subject: [PATCH 2/3] Added tests, fixed order by drain bug, and updated changelog --- ...smosCrossPartitionQueryExecutionContext.cs | 24 +- .../CosmosOrderByItemQueryExecutionContext.cs | 2 +- .../Query/MockPartitionResponse.cs | 17 +- .../Query/MockQueryFactory.cs | 77 ++++-- .../Query/QueryPipelineMockTests.cs | 229 +++++++++++++++++- .../Query/QueryResponseMessageFactory.cs | 28 ++- changelog.md | 2 +- 7 files changed, 329 insertions(+), 50 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index b955a3d19b..aa6841c603 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -354,14 +354,9 @@ public override async Task DrainAsync(int maxElements, Cancel { this.Stop(); - try - { - return this.FailureResponse; - } - finally - { - this.FailureResponse = null; - } + QueryResponseCore failure = this.FailureResponse.Value; + this.FailureResponse = null; + return failure; } // Drain the results. If there is no results and a failure then return the failure. @@ -369,15 +364,10 @@ public override async Task DrainAsync(int maxElements, Cancel if ((results == null || results.Count == 0) && this.FailureResponse != null) { this.Stop(); - - try - { - return this.FailureResponse; - } - finally - { - this.FailureResponse = null; - } + QueryResponseCore failure = this.FailureResponse.Value; + this.FailureResponse = null; + return failure; + } string continuation = this.ContinuationToken; diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs index 898f945fa8..159347ecfe 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs @@ -106,7 +106,7 @@ protected override string ContinuationToken // With this information we have captured the progress for all partitions in a single continuation token. get { - if (this.IsDone) + if (this.IsDone || this.FailureResponse != null) { return null; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockPartitionResponse.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockPartitionResponse.cs index fc85fc4bdf..de391ff1e4 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockPartitionResponse.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockPartitionResponse.cs @@ -15,6 +15,8 @@ namespace Microsoft.Azure.Cosmos.Tests /// internal class MockPartitionResponse { + internal const int MessageWithToManyRequestFailure = -1; + public PartitionKeyRange PartitionKeyRange { get; set; } /// @@ -24,6 +26,7 @@ internal class MockPartitionResponse /// /// /// Empty int[] represent an empty page + /// -1 represents a 429 failure /// public List MessagesWithItemIndex { get; set; } = new List(); @@ -46,15 +49,18 @@ public int GetTotalItemCount() int totalItemCount = 0; if(this.MessagesWithItemIndex != null) { - foreach (var message in this.MessagesWithItemIndex) + foreach (int[] message in this.MessagesWithItemIndex) { - totalItemCount += message.Length; + if (!IsFailurePage(message)) + { + totalItemCount += message.Length; + } } } if(this.Split != null) { - foreach(var partitionResponse in this.Split) + foreach(MockPartitionResponse partitionResponse in this.Split) { totalItemCount += partitionResponse.GetTotalItemCount(); } @@ -62,5 +68,10 @@ public int GetTotalItemCount() return totalItemCount; } + + public static bool IsFailurePage(int[] page) + { + return page != null && page.Length == 1 && page[0] < 0; + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs index 4f61548646..892069b92c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/MockQueryFactory.cs @@ -17,6 +17,7 @@ namespace Microsoft.Azure.Cosmos.Tests internal static class MockQueryFactory { public static readonly int[] EmptyPage = new int[] { }; + public static readonly int[] FailureToManyRequests = new int[] { MockPartitionResponse.MessageWithToManyRequestFailure }; public static readonly string DefaultDatabaseRid = ResourceId.NewDatabaseId(3810641).ToString(); public static readonly string DefaultCollectionRid = ResourceId.NewDocumentCollectionId(DefaultDatabaseRid, 1376573569).ToString(); public static readonly SqlQuerySpec DefaultQuerySpec = new SqlQuerySpec("SELECT * FROM C "); @@ -92,37 +93,53 @@ private static IList GenerateAndMockResponseHelper( string previousContinuationToken = initContinuationToken; // Loop through each message inside the partition - List messages = partitionAndMessages.MessagesWithItemIndex; - int messagesCount = messages == null ? 0 : messages.Count; - int lastMessageIndex = messagesCount - 1; - for (int i = 0; i < messagesCount; i++) + List pages = partitionAndMessages.MessagesWithItemIndex; + int pagesCount = pages == null ? 0 : pages.Count; + int lastPageIndex = pagesCount - 1; + for (int i = 0; i < pagesCount; i++) { - int[] message = partitionAndMessages.MessagesWithItemIndex[i]; - + int[] itemsInPage = partitionAndMessages.MessagesWithItemIndex[i]; string newContinuationToken = null; + QueryResponseCore queryResponse; - List currentPageItems = new List(); - // Null represents an empty page - if (message != null) + bool isFailureResponse = itemsInPage.Length == 1 && itemsInPage[0] < 0; + if (isFailureResponse) { - foreach (int itemPosition in message) + if(itemsInPage[0] == MockPartitionResponse.MessageWithToManyRequestFailure) + { + queryResponse = QueryResponseMessageFactory.CreateFailureToManyRequestResponse(); + } + else { - currentPageItems.Add(allItemsOrdered[itemPosition]); + throw new ArgumentException($"Unknown mocked failure response {itemsInPage[0]}"); } } - - // Last message should have null continuation token - // Split means it's not the last message for this PK range - if (i != lastMessageIndex || partitionAndMessages.HasSplit) + else { - newContinuationToken = Guid.NewGuid().ToString(); - } + List currentPageItems = new List(); + + // Null represents an empty page. Page with a single negative value represents a failure response + if (itemsInPage != null) + { + foreach (int itemPosition in itemsInPage) + { + currentPageItems.Add(allItemsOrdered[itemPosition]); + } + } + + // Last message should have null continuation token + // Split means it's not the last message for this PK range + if (i != lastPageIndex || partitionAndMessages.HasSplit) + { + newContinuationToken = Guid.NewGuid().ToString(); + } - QueryResponseCore queryResponse = QueryResponseMessageFactory.CreateQueryResponse( - currentPageItems, - isOrderByQuery, - newContinuationToken, - containerRid); + queryResponse = QueryResponseMessageFactory.CreateQueryResponse( + currentPageItems, + isOrderByQuery, + newContinuationToken, + containerRid); + } mockQueryClient.Setup(x => x.ExecuteItemQueryAsync( @@ -260,6 +277,22 @@ public static MockPartitionResponse[] CreateDefaultSplit( }; } + public static List GetFailureScenarios() + { + return new List + { + CreateDefaultResponse( + MockQueryFactory.FailureToManyRequests), + CreateDefaultResponse( + new int[] { 0 }, + MockQueryFactory.FailureToManyRequests), + CreateDefaultResponse( + new int[] { 0 }, + new int[] { 1 }, + MockQueryFactory.FailureToManyRequests), + }; + } + public static List GetSplitScenarios() { List allSplitScenario = new List(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs index cb87932e17..3dccb678c2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPipelineMockTests.cs @@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System; using System.Collections.Generic; using System.Linq; + using System.Net; using System.Text; using System.Threading; using System.Threading.Tasks; @@ -104,6 +105,102 @@ public async Task TestCosmosCrossPartitionQueryExecutionContextWithEmptyPagesAnd } } + + [TestMethod] + [DataRow(true)] + [DataRow(false)] + public async Task TestCosmosCrossPartitionQueryExecutionContextWithFailuresAsync(bool createInitialContinuationToken) + { + int maxPageSize = 5; + + List mockResponsesScenario = MockQueryFactory.GetFailureScenarios(); + foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) + { + string initialContinuationToken = null; + string fullConitnuationToken = null; + if (createInitialContinuationToken) + { + initialContinuationToken = " - RID:02FYAIvUH1kCAAAAAAAAAA ==#RT:1#TRC:1"; + CompositeContinuationToken compositeContinuation = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: MockQueryFactory.DefaultPartitionKeyRange.MinInclusive, + max: MockQueryFactory.DefaultPartitionKeyRange.MaxExclusive, + isMaxInclusive: false, + isMinInclusive: true), + Token = initialContinuationToken + }; + + fullConitnuationToken = JsonConvert.SerializeObject(new CompositeContinuationToken[] { compositeContinuation }); + } + + Mock mockQueryClient = new Mock(); + IList allItems = MockQueryFactory.GenerateAndMockResponse( + mockQueryClient, + isOrderByQuery: false, + sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, + containerRid: MockQueryFactory.DefaultCollectionRid, + initContinuationToken: initialContinuationToken, + maxPageSize: maxPageSize, + mockResponseForSinglePartition: mockResponse, + cancellationTokenForMocks: this.cancellationToken); + + CosmosQueryContext context = MockQueryFactory.CreateContext( + mockQueryClient.Object); + + CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, + collectionRid: MockQueryFactory.DefaultCollectionRid, + partitionedQueryExecutionInfo: new PartitionedQueryExecutionInfo() { QueryInfo = new QueryInfo() }, + partitionKeyRanges: new List() { MockQueryFactory.DefaultPartitionKeyRange }, + initialPageSize: maxPageSize, + maxConcurrency: null, + maxItemCount: maxPageSize, + maxBufferedItemCount: null); + + CosmosParallelItemQueryExecutionContext executionContext = await CosmosParallelItemQueryExecutionContext.CreateAsync( + context, + initParams, + fullConitnuationToken, + this.cancellationToken); + + // Read all the pages from both splits + List itemsRead = new List(); + Assert.IsTrue(!executionContext.IsDone); + + QueryResponseCore? failure = null; + while (!executionContext.IsDone) + { + QueryResponseCore queryResponse = await executionContext.DrainAsync(maxPageSize, this.cancellationToken); + if (queryResponse.IsSuccess) + { + string responseContinuationToken = queryResponse.ContinuationToken; + foreach (CosmosElement element in queryResponse.CosmosElements) + { + string jsonValue = element.ToString(); + ToDoItem item = JsonConvert.DeserializeObject(jsonValue); + itemsRead.Add(item); + } + } + else + { + Assert.IsNull(failure, "There should only be one error"); + failure = queryResponse; + } + } + + Assert.IsNotNull(failure); + Assert.AreEqual((HttpStatusCode)429, failure.Value.StatusCode); + Assert.IsNull(failure.Value.ErrorMessage); + + Assert.AreEqual(allItems.Count, itemsRead.Count); + List exepected = allItems.OrderBy(x => x.id).ToList(); + List actual = itemsRead.OrderBy(x => x.id).ToList(); + + CollectionAssert.AreEqual(exepected, actual, new ToDoItemComparer()); + } + } + [TestMethod] [DataRow(true)] [DataRow(false)] @@ -152,7 +249,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs fullConitnuationToken = JsonConvert.SerializeObject(new OrderByContinuationToken[] { orderByContinuationToken }); } - + IList allItems = MockQueryFactory.GenerateAndMockResponse( mockQueryClient, isOrderByQuery: true, @@ -165,7 +262,7 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs // Order by drains the partitions until it finds an item // If there are no items then it's not possible to have a continuation token - if(allItems.Count == 0 && createInitialContinuationToken) + if (allItems.Count == 0 && createInitialContinuationToken) { continue; } @@ -226,5 +323,133 @@ public async Task TestCosmosOrderByQueryExecutionContextWithEmptyPagesAndSplitAs CollectionAssert.AreEqual(allItems.ToList(), itemsRead, new ToDoItemComparer()); } } + + [TestMethod] + [DataRow(true)] + [DataRow(false)] + public async Task TestCosmosOrderByQueryExecutionContextWithFailurePageAsync(bool createInitialContinuationToken) + { + int maxPageSize = 5; + + List mockResponsesScenario = MockQueryFactory.GetFailureScenarios(); + + Mock mockQueryClient = new Mock(); + foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario) + { + string initialContinuationToken = null; + string fullConitnuationToken = null; + if (createInitialContinuationToken) + { + ToDoItem itemToRepresentPreviousQuery = ToDoItem.CreateItems( + 1, + "itemToRepresentPreviousQuery", + MockQueryFactory.DefaultCollectionRid).First(); + + initialContinuationToken = $" - RID:{itemToRepresentPreviousQuery._rid} ==#RT:1#TRC:1"; + CompositeContinuationToken compositeContinuation = new CompositeContinuationToken() + { + Range = new Documents.Routing.Range( + min: MockQueryFactory.DefaultPartitionKeyRange.MinInclusive, + max: MockQueryFactory.DefaultPartitionKeyRange.MaxExclusive, + isMaxInclusive: false, + isMinInclusive: true), + Token = initialContinuationToken + }; + + List orderByItems = new List() + { + new OrderByItem(CosmosObject.CreateFromBuffer(Encoding.UTF8.GetBytes("{\"item\":\"2c4ce711-13c3-4c93-817c-49287b71b6c3\"}"))) + }; + + OrderByContinuationToken orderByContinuationToken = new OrderByContinuationToken( + queryClient: mockQueryClient.Object, + compositeContinuationToken: compositeContinuation, + orderByItems: orderByItems, + rid: itemToRepresentPreviousQuery._rid, + skipCount: 0, + filter: null); + + fullConitnuationToken = JsonConvert.SerializeObject(new OrderByContinuationToken[] { orderByContinuationToken }); + } + + + IList allItems = MockQueryFactory.GenerateAndMockResponse( + mockQueryClient, + isOrderByQuery: true, + sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, + containerRid: MockQueryFactory.DefaultCollectionRid, + initContinuationToken: initialContinuationToken, + maxPageSize: maxPageSize, + mockResponseForSinglePartition: mockResponse, + cancellationTokenForMocks: this.cancellationToken); + + // Order by drains the partitions until it finds an item + // If there are no items then it's not possible to have a continuation token + if (allItems.Count == 0 && createInitialContinuationToken) + { + continue; + } + + CosmosQueryContext context = MockQueryFactory.CreateContext( + mockQueryClient.Object); + + QueryInfo queryInfo = new QueryInfo() + { + OrderBy = new SortOrder[] { SortOrder.Ascending }, + OrderByExpressions = new string[] { "id" } + }; + + CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams initParams = new CosmosCrossPartitionQueryExecutionContext.CrossPartitionInitParams( + sqlQuerySpec: MockQueryFactory.DefaultQuerySpec, + collectionRid: MockQueryFactory.DefaultCollectionRid, + partitionedQueryExecutionInfo: new PartitionedQueryExecutionInfo() { QueryInfo = queryInfo }, + partitionKeyRanges: new List() { MockQueryFactory.DefaultPartitionKeyRange }, + initialPageSize: maxPageSize, + maxConcurrency: null, + maxItemCount: maxPageSize, + maxBufferedItemCount: null); + + CosmosOrderByItemQueryExecutionContext executionContext = await CosmosOrderByItemQueryExecutionContext.CreateAsync( + context, + initParams, + fullConitnuationToken, + this.cancellationToken); + + Assert.IsTrue(!executionContext.IsDone); + + // Read all the pages from both splits + List itemsRead = new List(); + QueryResponseCore? failure = null; + while (!executionContext.IsDone) + { + QueryResponseCore queryResponse = await executionContext.DrainAsync( + maxPageSize, + this.cancellationToken); + if (queryResponse.IsSuccess) + { + string responseContinuationToken = queryResponse.ContinuationToken; + foreach (CosmosElement element in queryResponse.CosmosElements) + { + string jsonValue = element.ToString(); + ToDoItem item = JsonConvert.DeserializeObject(jsonValue); + itemsRead.Add(item); + } + } + else + { + Assert.IsNull(failure, "There should only be one error"); + failure = queryResponse; + } + } + + Assert.IsNotNull(failure); + Assert.AreEqual((HttpStatusCode)429, failure.Value.StatusCode); + Assert.IsNull(failure.Value.ErrorMessage); + + Assert.AreEqual(allItems.Count, itemsRead.Count); + + CollectionAssert.AreEqual(allItems.ToList(), itemsRead, new ToDoItemComparer()); + } + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryResponseMessageFactory.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryResponseMessageFactory.cs index b477e92acd..54bcd99047 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryResponseMessageFactory.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryResponseMessageFactory.cs @@ -99,12 +99,15 @@ public static QueryResponseCore CreateQueryResponse( return message; } - public static QueryResponseCore CreateSplitResponse(string collectionRid) + public static QueryResponseCore CreateFailureResponse( + HttpStatusCode httpStatusCode, + SubStatusCodes subStatusCodes, + string errorMessage) { QueryResponseCore splitResponse = QueryResponseCore.CreateFailure( - statusCode: HttpStatusCode.Gone, - subStatusCodes: SubStatusCodes.PartitionKeyRangeGone, - errorMessage: "Partition split error", + statusCode: httpStatusCode, + subStatusCodes: subStatusCodes, + errorMessage: errorMessage, requestCharge: 10.4, activityId: Guid.NewGuid().ToString(), queryMetricsText: null, @@ -113,6 +116,23 @@ public static QueryResponseCore CreateSplitResponse(string collectionRid) return splitResponse; } + public static QueryResponseCore CreateFailureToManyRequestResponse() + { + // 429 do not have an error message + return CreateFailureResponse( + (HttpStatusCode)429, + SubStatusCodes.Unknown, + null); + } + + public static QueryResponseCore CreateSplitResponse(string collectionRid) + { + return CreateFailureResponse( + HttpStatusCode.Gone, + SubStatusCodes.PartitionKeyRangeGone, + "Partition split error"); + } + private static MemoryStream SerializeForOrderByQuery(IList items) { OrderByReturnStructure[] payload = items.Select(item => new OrderByReturnStructure() diff --git a/changelog.md b/changelog.md index cc85828149..76b170c0b9 100644 --- a/changelog.md +++ b/changelog.md @@ -16,7 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - [#905](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/909) Fixed linq camel case bug - +- [#927](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/927) Fixed query returning partial results instead of error ## [3.3.1](https://www.nuget.org/packages/Microsoft.Azure.Cosmos/3.3.1) - 2019-10-11 From 33748f7dae188b1f9f253e87f0ecfb7b11974c53 Mon Sep 17 00:00:00 2001 From: Jake Willey Date: Tue, 22 Oct 2019 10:04:03 -0700 Subject: [PATCH 3/3] Updated the failure handling --- .../CosmosCrossPartitionQueryExecutionContext.cs | 2 +- .../CosmosOrderByItemQueryExecutionContext.cs | 10 ++++------ .../CosmosParallelItemQueryExecutionContext.cs | 2 +- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs index aa6841c603..6fe01206db 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosCrossPartitionQueryExecutionContext.cs @@ -336,7 +336,7 @@ protected async Task MoveNextHelperAsync(ItemProducerTree itemProducerTree this.FailureResponse = moveNextResponse.failureResponse; } - return !moveNextResponse.successfullyMovedNext; + return moveNextResponse.successfullyMovedNext; } /// diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs index 159347ecfe..8ed3b4101e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosOrderByItemQueryExecutionContext.cs @@ -106,7 +106,7 @@ protected override string ContinuationToken // With this information we have captured the progress for all partitions in a single continuation token. get { - if (this.IsDone || this.FailureResponse != null) + if (this.IsDone) { return null; } @@ -209,7 +209,8 @@ public override async Task> InternalDrainAsync(int //// 2) always come before where j < k List results = new List(); - while (!this.IsDone && results.Count < maxElements) + bool isSuccessToMoveNext = true; + while (!this.IsDone && results.Count < maxElements && isSuccessToMoveNext) { // Only drain from the highest priority document producer // We need to pop and push back the document producer tree, since the priority changes according to the sort order. @@ -234,10 +235,7 @@ public override async Task> InternalDrainAsync(int this.previousRid = orderByQueryResult.Rid; - if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken)) - { - break; - } + isSuccessToMoveNext = await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken); this.PushCurrentItemProducerTree(currentItemProducerTree); } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs index 682c029a2e..270b587958 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosParallelItemQueryExecutionContext.cs @@ -155,7 +155,7 @@ public override async Task> InternalDrainAsync(int for (int i = 0; i < Math.Min(itemsLeftInCurrentPage, maxElements); i++) { results.Add(currentItemProducerTree.Current); - if (await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken)) + if (!await this.MoveNextHelperAsync(currentItemProducerTree, cancellationToken)) { break; }