private static TryCatch <CrossFeedRangeState <QueryState> > MonadicExtractState(
CosmosElement continuationToken,
IReadOnlyList <FeedRangeEpk> ranges)
{
if (continuationToken == null)
{
// Full fan out to the ranges with null continuations
CrossFeedRangeState <QueryState> fullFanOutState = new CrossFeedRangeState <QueryState>(ranges.Select(range => new FeedRangeState <QueryState>(range, (QueryState)null)).ToArray());
return(TryCatch <CrossFeedRangeState <QueryState> > .FromResult(fullFanOutState));
}
if (!(continuationToken is CosmosArray parallelContinuationTokenListRaw))
{
return(TryCatch <CrossFeedRangeState <QueryState> > .FromException(
new MalformedContinuationTokenException(
$"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")));
}
if (parallelContinuationTokenListRaw.Count == 0)
{
return(TryCatch <CrossFeedRangeState <QueryState> > .FromException(
new MalformedContinuationTokenException(
$"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")));
}
List <ParallelContinuationToken> parallelContinuationTokens = new List <ParallelContinuationToken>();
foreach (CosmosElement parallelContinuationTokenRaw in parallelContinuationTokenListRaw)
{
TryCatch <ParallelContinuationToken> tryCreateParallelContinuationToken = ParallelContinuationToken.TryCreateFromCosmosElement(parallelContinuationTokenRaw);
if (tryCreateParallelContinuationToken.Failed)
{
return(TryCatch <CrossFeedRangeState <QueryState> > .FromException(
tryCreateParallelContinuationToken.Exception));
}
parallelContinuationTokens.Add(tryCreateParallelContinuationToken.Result);
}
TryCatch <PartitionMapping <ParallelContinuationToken> > partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping(
ranges,
parallelContinuationTokens);
if (partitionMappingMonad.Failed)
{
return(TryCatch <CrossFeedRangeState <QueryState> > .FromException(
partitionMappingMonad.Exception));
}
PartitionMapping <ParallelContinuationToken> partitionMapping = partitionMappingMonad.Result;
List <FeedRangeState <QueryState> > feedRangeStates = new List <FeedRangeState <QueryState> >();
List <IReadOnlyDictionary <FeedRangeEpk, ParallelContinuationToken> > rangesToInitialize = new List <IReadOnlyDictionary <FeedRangeEpk, ParallelContinuationToken> >()
{
// Skip all the partitions left of the target range, since they have already been drained fully.
partitionMapping.TargetMapping,
partitionMapping.MappingRightOfTarget,
};
foreach (IReadOnlyDictionary <FeedRangeEpk, ParallelContinuationToken> rangeToInitalize in rangesToInitialize)
{
foreach (KeyValuePair <FeedRangeEpk, ParallelContinuationToken> kvp in rangeToInitalize)
{
FeedRangeState <QueryState> feedRangeState = new FeedRangeState <QueryState>(kvp.Key, kvp.Value?.Token != null ? new QueryState(CosmosString.Create(kvp.Value.Token)) : null);
feedRangeStates.Add(feedRangeState);
}
}
CrossFeedRangeState <QueryState> crossPartitionState = new CrossFeedRangeState <QueryState>(feedRangeStates.ToArray());
return(TryCatch <CrossFeedRangeState <QueryState> > .FromResult(crossPartitionState));
}
private static void RunTryGetInitializationInfo(
IReadOnlyDictionary <PartitionKeyRange, IPartitionedToken> expectedMappingLeftPartitions,
IReadOnlyDictionary <PartitionKeyRange, IPartitionedToken> expectedMappingTargetPartition,
IReadOnlyDictionary <PartitionKeyRange, IPartitionedToken> expectedMappingRightPartitions,
IEnumerable <PartitionKeyRange> partitionKeyRanges,
IEnumerable <IPartitionedToken> partitionedTokens)
{
TryCatch <PartitionMapping <IPartitionedToken> > tryGetInitializationInfo = PartitionMapper.MonadicGetPartitionMapping <IPartitionedToken>(
partitionKeyRanges.OrderBy(x => Guid.NewGuid()).ToArray(),
partitionedTokens.OrderBy(x => Guid.NewGuid()).ToList());
Assert.IsTrue(tryGetInitializationInfo.Succeeded);
PartitionMapping <IPartitionedToken> partitionMapping = tryGetInitializationInfo.Result;
AssertPartitionMappingAreEqual(expectedMappingLeftPartitions, partitionMapping.PartitionsLeftOfTarget);
AssertPartitionMappingAreEqual(expectedMappingTargetPartition, partitionMapping.TargetPartition);
AssertPartitionMappingAreEqual(expectedMappingRightPartitions, partitionMapping.PartitionsRightOfTarget);
}