public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// In order to maintain the continuation token for the user we must drain with a few constraints
// 1) We fully drain from the left most partition before moving on to the next partition
// 2) We drain only full pages from the document producer so we aren't left with a partial page
// otherwise we would need to add to the continuation token how many items to skip over on that page.
// Only drain from the leftmost (current) document producer tree
ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree();
List <CosmosElement> results = new List <CosmosElement>();
try
{
(bool gotNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken);
if (failureResponse != null)
{
return(failureResponse.Value);
}
if (gotNextPage)
{
int itemsLeftInCurrentPage = currentItemProducerTree.ItemsLeftInCurrentPage;
// Only drain full pages or less if this is a top query.
currentItemProducerTree.TryMoveNextDocumentWithinPage();
int numberOfItemsToDrain = Math.Min(itemsLeftInCurrentPage, maxElements);
for (int i = 0; i < numberOfItemsToDrain; i++)
{
results.Add(currentItemProducerTree.Current);
currentItemProducerTree.TryMoveNextDocumentWithinPage();
}
}
}
finally
{
this.PushCurrentItemProducerTree(currentItemProducerTree);
}
return(QueryResponseCore.CreateSuccess(
result: results,
requestCharge: this.requestChargeTracker.GetAndResetCharge(),
activityId: null,
responseLengthBytes: this.GetAndResetResponseLengthBytes(),
disallowContinuationTokenMessage: null,
continuationToken: this.ContinuationToken,
diagnostics: this.GetAndResetDiagnostics()));
}
/// <summary>
/// Initializes cross partition query execution context by initializing the necessary document producers.
/// </summary>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="partitionKeyRanges">The partitions to target.</param>
/// <param name="initialPageSize">The page size to start the document producers off with.</param>
/// <param name="querySpecForInit">The query specification for the rewritten query.</param>
/// <param name="targetRangeToContinuationMap">Map from partition to it's corresponding continuation token.</param>
/// <param name="deferFirstPage">Whether or not we should defer the fetch of the first page from each partition.</param>
/// <param name="filter">The filter to inject in the predicate.</param>
/// <param name="tryFilterAsync">The callback used to filter each partition.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task to await on.</returns>
protected async Task <TryCatch> TryInitializeAsync(
string collectionRid,
IReadOnlyList <PartitionKeyRange> partitionKeyRanges,
int initialPageSize,
SqlQuerySpec querySpecForInit,
IReadOnlyDictionary <string, string> targetRangeToContinuationMap,
bool deferFirstPage,
string filter,
Func <ItemProducerTree, Task <TryCatch> > tryFilterAsync,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
List <ItemProducerTree> itemProducerTrees = new List <ItemProducerTree>();
foreach (PartitionKeyRange partitionKeyRange in partitionKeyRanges)
{
string initialContinuationToken;
if (targetRangeToContinuationMap != null)
{
if (!targetRangeToContinuationMap.TryGetValue(partitionKeyRange.Id, out initialContinuationToken))
{
initialContinuationToken = null;
}
}
else
{
initialContinuationToken = null;
}
ItemProducerTree itemProducerTree = new ItemProducerTree(
this.queryContext,
querySpecForInit,
partitionKeyRange,
this.OnItemProducerTreeCompleteFetching,
this.itemProducerForest.Comparer as IComparer <ItemProducerTree>,
this.equalityComparer,
this.testSettings,
deferFirstPage,
collectionRid,
initialPageSize,
initialContinuationToken)
{
Filter = filter
};
// Prefetch if necessary, and populate consume queue.
if (this.CanPrefetch)
{
this.TryScheduleFetch(itemProducerTree);
}
itemProducerTrees.Add(itemProducerTree);
}
// Using loop fission so that we can load the document producers in parallel
foreach (ItemProducerTree itemProducerTree in itemProducerTrees)
{
if (!deferFirstPage)
{
while (true)
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationToken);
if (failureResponse.HasValue)
{
return(TryCatch.FromException(
new CosmosException(
statusCode: failureResponse.Value.StatusCode,
subStatusCode: (int)failureResponse.Value.SubStatusCode.GetValueOrDefault(0),
message: failureResponse.Value.ErrorMessage,
activityId: failureResponse.Value.ActivityId,
requestCharge: failureResponse.Value.RequestCharge)));
}
if (!movedToNextPage)
{
break;
}
if (itemProducerTree.IsAtBeginningOfPage)
{
break;
}
if (itemProducerTree.TryMoveNextDocumentWithinPage())
{
break;
}
}
}
if (tryFilterAsync != null)
{
TryCatch tryFilter = await tryFilterAsync(itemProducerTree);
if (!tryFilter.Succeeded)
{
return(tryFilter);
}
}
this.itemProducerForest.Enqueue(itemProducerTree);
}
return(TryCatch.FromResult());
}