public async Task Tracing()
{
List <CosmosObject> documents = new List <CosmosObject>();
for (int i = 0; i < 250; i++)
{
documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}"));
}
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(documents);
IQueryPipelineStage pipelineStage = await CreatePipelineAsync(documentContainer, "SELECT * FROM c", pageSize : 10);
Trace rootTrace;
int numTraces = (await documentContainer.GetFeedRangesAsync(NoOpTrace.Singleton, default)).Count;
using (rootTrace = Trace.GetRootTrace("Cross Partition Query"))
{
while (await pipelineStage.MoveNextAsync(rootTrace))
{
TryCatch <QueryPage> tryGetQueryPage = pipelineStage.Current;
tryGetQueryPage.ThrowIfFailed();
numTraces++;
}
}
Assert.AreEqual(numTraces, rootTrace.Children.Count);
}
public async Task Tracing()
{
List <CosmosObject> documents = new List <CosmosObject>();
for (int i = 0; i < 250; i++)
{
documents.Add(CosmosObject.Parse($"{{\"pk\" : {i} }}"));
}
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(documents);
IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, "SELECT * FROM c", pageSize: 10);
Trace rootTrace;
int numTraces = 1;
using (rootTrace = Trace.GetRootTrace("Cross Partition Query"))
{
while (await pipelineStage.MoveNextAsync(rootTrace))
{
TryCatch <QueryPage> tryGetQueryPage = pipelineStage.Current;
tryGetQueryPage.ThrowIfFailed();
numTraces++;
}
}
string traceString = TraceWriter.TraceToText(rootTrace);
Console.WriteLine(traceString);
Assert.AreEqual(numTraces, rootTrace.Children.Count);
}
public async ValueTask <bool> MoveNextAsync()
{
if (!await this.currentQueryPipelineStage.MoveNextAsync())
{
return(false);
}
TryCatch <QueryPage> tryGetSourcePage = this.currentQueryPipelineStage.Current;
this.Current = tryGetSourcePage;
if (tryGetSourcePage.Failed)
{
Exception exception = tryGetSourcePage.InnerMostException;
bool shouldRetry = (exception is CosmosException cosmosException) &&
(cosmosException.StatusCode == HttpStatusCode.Gone) &&
(cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.NameCacheIsStale) &&
!this.alreadyRetried;
if (shouldRetry)
{
await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync(
this.cosmosQueryContext.ResourceLink,
default);
this.alreadyRetried = true;
await this.currentQueryPipelineStage.DisposeAsync();
this.currentQueryPipelineStage = this.queryPipelineStageFactory();
return(await this.MoveNextAsync());
}
}
return(true);
}
private ComputeDistinctQueryPipelineStage(
DistinctMap distinctMap,
IQueryPipelineStage source,
CancellationToken cancellationToken)
: base(distinctMap, source, cancellationToken)
{
}
protected DistinctQueryPipelineStage(
DistinctMap distinctMap,
IQueryPipelineStage source,
CancellationToken cancellationToken)
: base(source, cancellationToken)
{
this.distinctMap = distinctMap ?? throw new ArgumentNullException(nameof(distinctMap));
}
private ComputeTakeQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
int takeCount)
: base(source, cancellationToken, takeCount)
{
// Work is done in the base class.
}
private ClientSkipQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
long skipCount)
: base(source, cancellationToken, skipCount)
{
// Work is done in base constructor.
}
private ComputeGroupByQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
GroupingTable groupingTable,
int pageSize)
: base(source, cancellationToken, groupingTable, pageSize)
{
}
public NameCacheStaleRetryQueryPipelineStage(
CosmosQueryContext cosmosQueryContext,
Func <IQueryPipelineStage> queryPipelineStageFactory)
{
this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext));
this.queryPipelineStageFactory = queryPipelineStageFactory ?? throw new ArgumentNullException(nameof(queryPipelineStageFactory));
this.currentQueryPipelineStage = queryPipelineStageFactory();
}
protected TakeQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
int takeCount)
: base(source, cancellationToken)
{
this.takeCount = takeCount;
}
public async Task BasicTests()
{
IReadOnlyList <TestCase> testCases = new List <TestCase>()
{
MakeTest(
input: Enumerable
.Repeat(EmptyPagePipelineStage.PageType.Empty, 2000)
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 1)),
expected: Enumerable.Repeat(true, 1)),
MakeTest(
input: Enumerable
.Repeat(EmptyPagePipelineStage.PageType.Empty, 100)
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 5))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Empty, 27))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 3))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Empty, 32))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 3)),
expected: Enumerable.Repeat(true, 11)),
MakeTest(
input: Enumerable
.Repeat(EmptyPagePipelineStage.PageType.Empty, 100)
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 5))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Empty, 27))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Error, 3))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Empty, 32))
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 3)),
expected: Enumerable.Repeat(true, 5)
.Concat(Enumerable.Repeat(false, 3))
.Concat(Enumerable.Repeat(true, 3))),
MakeTest(
input: Enumerable.Repeat(EmptyPagePipelineStage.PageType.NonEmpty, 500),
expected: Enumerable.Repeat(true, 500)),
MakeTest(
input: Enumerable.Repeat(EmptyPagePipelineStage.PageType.Error, 500),
expected: Enumerable.Repeat(false, 500))
};
foreach (TestCase testCase in testCases)
{
await using IQueryPipelineStage pipeline = CreatePipeline(testCase.Input);
for (int index = 0; index < testCase.Expected.Count; ++index)
{
Assert.IsTrue(await pipeline.MoveNextAsync(NoOpTrace.Singleton));
if (testCase.Expected[index])
{
Assert.IsTrue(pipeline.Current.Succeeded);
Assert.AreEqual(1, pipeline.Current.Result.Documents.Count);
Assert.AreEqual("42", pipeline.Current.Result.Documents[0].ToString());
}
else
{
Assert.IsTrue(pipeline.Current.Failed);
}
}
}
}
private ClientTakeQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
int takeCount,
TakeEnum takeEnum)
: base(source, cancellationToken, takeCount)
{
this.takeEnum = takeEnum;
}
private ComputeAggregateQueryPipelineStage(
IQueryPipelineStage source,
SingleGroupAggregator singleGroupAggregator,
bool isValueAggregateQuery,
CancellationToken cancellationToken)
: base(source, singleGroupAggregator, isValueAggregateQuery, cancellationToken)
{
// all the work is done in the base constructor.
}
private ComputeDCountQueryPipelineStage(
IQueryPipelineStage source,
long count,
DCountInfo info,
CancellationToken cancellationToken)
: base(source, count, info, cancellationToken)
{
// all the work is done in the base constructor.
}
protected GroupByQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
GroupingTable groupingTable,
int pageSize)
: base(source, cancellationToken)
{
this.groupingTable = groupingTable ?? throw new ArgumentNullException(nameof(groupingTable));
this.pageSize = pageSize;
}
/// <summary>
/// Initializes a new instance of the AggregateDocumentQueryExecutionComponent class.
/// </summary>
/// <param name="source">The source component that will supply the local aggregates from multiple continuations and partitions.</param>
/// <param name="singleGroupAggregator">The single group aggregator that we will feed results into.</param>
/// <param name="isValueQuery">Whether or not the query has the 'VALUE' keyword.</param>
/// <param name="cancellationToken">The cancellation token for cooperative yeilding.</param>
/// <remarks>This constructor is private since there is some async initialization that needs to happen in CreateAsync().</remarks>
public AggregateQueryPipelineStage(
IQueryPipelineStage source,
SingleGroupAggregator singleGroupAggregator,
bool isValueQuery,
CancellationToken cancellationToken)
: base(source, cancellationToken)
{
this.singleGroupAggregator = singleGroupAggregator ?? throw new ArgumentNullException(nameof(singleGroupAggregator));
this.isValueQuery = isValueQuery;
}
/// <summary>
/// Initializes a new instance of the DCountQueryPipelineStage class.
/// </summary>
/// <param name="source">The source component that will supply the local aggregates from multiple continuations and partitions.</param>
/// <param name="count">The actual dcount that will be reported.</param>
/// <param name="info">Metadata about the original dcount query that is elided in the rewritten query</param>
/// <param name="cancellationToken">The cancellation token for cooperative yeilding.</param>
/// <remarks>This constructor is private since there is some async initialization that needs to happen in CreateAsync().</remarks>
public DCountQueryPipelineStage(
IQueryPipelineStage source,
long count,
DCountInfo info,
CancellationToken cancellationToken)
: base(source, cancellationToken)
{
this.count = count;
this.info = info;
}
public async Task StackOverflowTest()
{
await using IQueryPipelineStage pipeline = CreatePipeline(Enumerable
.Repeat(EmptyPagePipelineStage.PageType.Empty, 2000)
.Concat(Enumerable.Repeat(EmptyPagePipelineStage.PageType.Error, 1))
.ToList());
bool hasNext = await pipeline.MoveNextAsync(NoOpTrace.Singleton);
Assert.IsTrue(hasNext);
TryCatch <QueryPage> result = pipeline.Current;
Assert.IsFalse(result.Succeeded);
}
protected SkipQueryPipelineStage(
IQueryPipelineStage source,
CancellationToken cancellationToken,
long skipCount)
: base(source, cancellationToken)
{
if (skipCount > int.MaxValue)
{
throw new ArgumentOutOfRangeException(nameof(skipCount));
}
this.skipCount = (int)skipCount;
}
private ClientDistinctQueryPipelineStage(
DistinctQueryType distinctQueryType,
DistinctMap distinctMap,
IQueryPipelineStage source,
CancellationToken cancellationToken)
: base(distinctMap, source, cancellationToken)
{
if ((distinctQueryType != DistinctQueryType.Unordered) && (distinctQueryType != DistinctQueryType.Ordered))
{
throw new ArgumentException($"Unknown {nameof(DistinctQueryType)}: {distinctQueryType}.");
}
this.distinctQueryType = distinctQueryType;
}
private QueryIterator(
CosmosQueryContextCore cosmosQueryContext,
IQueryPipelineStage cosmosQueryExecutionContext,
CosmosSerializationFormatOptions cosmosSerializationFormatOptions,
RequestOptions requestOptions,
CosmosClientContext clientContext)
{
this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext));
this.queryPipelineStage = cosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContext));
this.cosmosSerializationFormatOptions = cosmosSerializationFormatOptions;
this.requestOptions = requestOptions;
this.clientContext = clientContext ?? throw new ArgumentNullException(nameof(clientContext));
this.hasMoreResults = true;
}
private static async Task <List <CosmosElement> > DrainWithoutStateAsync(string query, IDocumentContainer documentContainer, int pageSize = 10)
{
IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, query, pageSize);
List <CosmosElement> elements = new List <CosmosElement>();
while (await pipelineStage.MoveNextAsync())
{
TryCatch <QueryPage> tryGetQueryPage = pipelineStage.Current;
tryGetQueryPage.ThrowIfFailed();
elements.AddRange(tryGetQueryPage.Result.Documents);
}
return(elements);
}
public async Task TestMerge()
{
List <CosmosObject> documents = Enumerable
.Range(0, 100)
.Select(x => CosmosObject.Parse($"{{\"pk\" : {x} }}"))
.ToList();
MergeTestUtil mergeTest = new MergeTestUtil();
mergeTest.DocumentContainer = await CreateDocumentContainerAsync(
documents : documents,
numPartitions : 2,
failureConfigs : new FlakyDocumentContainer.FailureConfigs(
inject429s: false,
injectEmptyPages: false,
shouldReturnFailure: mergeTest.ShouldReturnFailure));
string query = "SELECT * FROM c ORDER BY c._ts";
int pageSize = 10;
IQueryPipelineStage pipelineStage = CreatePipeline(mergeTest.DocumentContainer, query, pageSize);
List <CosmosElement> elements = new List <CosmosElement>();
int iteration = 0;
while (await pipelineStage.MoveNextAsync())
{
TryCatch <QueryPage> tryGetQueryPage = pipelineStage.Current;
tryGetQueryPage.ThrowIfFailed();
elements.AddRange(tryGetQueryPage.Result.Documents);
++iteration;
if (iteration == 1)
{
mergeTest.ShouldMerge = MergeTestUtil.TriState.Ready;
}
}
Assert.AreEqual(expected: documents.Count, actual: elements.Count);
}
public async Task QueryAsync()
{
int numItems = 100;
IDocumentContainer documentContainer = await this.CreateDocumentContainerAsync(numItems);
IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, "SELECT * FROM c", pageSize: 10);
Trace rootTrace;
int numChildren = 1; // One extra since we need to read one past the last user page to get the null continuation.
using (rootTrace = Trace.GetRootTrace("Cross Partition Query"))
{
while (await pipelineStage.MoveNextAsync(rootTrace))
{
numChildren++;
}
}
string traceString = TraceWriter.TraceToText(rootTrace);
Console.WriteLine(traceString);
Assert.AreEqual(numChildren, rootTrace.Children.Count);
}
protected QueryPipelineStageBase(IQueryPipelineStage inputStage, CancellationToken cancellationToken)
{
this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage));
this.cancellationToken = cancellationToken;
}
public CatchAllQueryPipelineStage(IQueryPipelineStage inputStage, CancellationToken cancellationToken)
: base(inputStage, cancellationToken)
{
}
public SkipEmptyPageQueryPipelineStage(IQueryPipelineStage inputStage, CancellationToken cancellationToken)
{
this.inputStage = inputStage ?? throw new ArgumentNullException(nameof(inputStage));
this.cancellationToken = cancellationToken;
}
public async Task ScenariosAsync()
{
List <Input> inputs = new List <Input>();
int startLineNumber;
int endLineNumber;
//----------------------------------------------------------------
// ReadFeed
//----------------------------------------------------------------
{
startLineNumber = GetLineNumber();
int numItems = 100;
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems);
CrossPartitionReadFeedAsyncEnumerator enumerator = CrossPartitionReadFeedAsyncEnumerator.Create(
documentContainer,
new CrossFeedRangeState <ReadFeedState>(ReadFeedCrossFeedRangeState.CreateFromBeginning().FeedRangeStates),
new ReadFeedPaginationOptions(pageSizeHint: 10),
cancellationToken: default);
int numChildren = 1; // One extra since we need to read one past the last user page to get the null continuation.
TraceForBaselineTesting rootTrace;
using (rootTrace = TraceForBaselineTesting.GetRootTrace())
{
while (await enumerator.MoveNextAsync(rootTrace))
{
numChildren++;
}
}
Assert.AreEqual(numChildren, rootTrace.Children.Count);
endLineNumber = GetLineNumber();
inputs.Add(new Input("ReadFeed", rootTrace, startLineNumber, endLineNumber));
}
//----------------------------------------------------------------
//----------------------------------------------------------------
// ChangeFeed
//----------------------------------------------------------------
{
startLineNumber = GetLineNumber();
int numItems = 100;
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems);
CrossPartitionChangeFeedAsyncEnumerator enumerator = CrossPartitionChangeFeedAsyncEnumerator.Create(
documentContainer,
new CrossFeedRangeState <ChangeFeedState>(
ChangeFeedCrossFeedRangeState.CreateFromBeginning().FeedRangeStates),
new ChangeFeedPaginationOptions(
ChangeFeedMode.Incremental,
pageSizeHint: int.MaxValue),
cancellationToken: default);
int numChildren = 0;
TraceForBaselineTesting rootTrace;
using (rootTrace = TraceForBaselineTesting.GetRootTrace())
{
while (await enumerator.MoveNextAsync(rootTrace))
{
numChildren++;
if (enumerator.Current.Result.Page is ChangeFeedNotModifiedPage)
{
break;
}
}
}
Assert.AreEqual(numChildren, rootTrace.Children.Count);
endLineNumber = GetLineNumber();
inputs.Add(new Input("ChangeFeed", rootTrace, startLineNumber, endLineNumber));
}
//----------------------------------------------------------------
//----------------------------------------------------------------
// Query
//----------------------------------------------------------------
{
startLineNumber = GetLineNumber();
int numItems = 100;
IDocumentContainer documentContainer = await CreateDocumentContainerAsync(numItems);
IQueryPipelineStage pipelineStage = CreatePipeline(documentContainer, "SELECT * FROM c", pageSize: 10);
TraceForBaselineTesting rootTrace;
int numChildren = 1; // One extra since we need to read one past the last user page to get the null continuation.
using (rootTrace = TraceForBaselineTesting.GetRootTrace())
{
while (await pipelineStage.MoveNextAsync(rootTrace))
{
numChildren++;
}
}
Assert.AreEqual(numChildren, rootTrace.Children.Count);
endLineNumber = GetLineNumber();
inputs.Add(new Input("Query", rootTrace, startLineNumber, endLineNumber));
}
//----------------------------------------------------------------
this.ExecuteTestSuite(inputs);
}
public async Task TestCosmosQueryPartitionKeyDefinition()
{
PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition();
QueryRequestOptions queryRequestOptions = new QueryRequestOptions
{
Properties = new Dictionary <string, object>()
{
{ "x-ms-query-partitionkey-definition", partitionKeyDefinition }
}
};
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(@"select * from t where t.something = 42 ");
bool allowNonValueAggregateQuery = true;
bool isContinuationExpected = true;
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
CancellationToken cancellationtoken = cancellationTokenSource.Token;
Mock <CosmosQueryClient> client = new Mock <CosmosQueryClient>();
string exceptionMessage = "Verified that the PartitionKeyDefinition was correctly set. Cancel the rest of the query";
client
.Setup(x => x.GetCachedContainerQueryPropertiesAsync(It.IsAny <string>(), It.IsAny <Cosmos.PartitionKey?>(), cancellationtoken))
.ReturnsAsync(new ContainerQueryProperties("mockContainer", null, partitionKeyDefinition));
client
.Setup(x => x.ByPassQueryParsing())
.Returns(false);
client
.Setup(x => x.TryGetPartitionedQueryExecutionInfoAsync(
It.IsAny <SqlQuerySpec>(),
It.IsAny <PartitionKeyDefinition>(),
It.IsAny <bool>(),
It.IsAny <bool>(),
It.IsAny <bool>(),
It.IsAny <bool>(),
It.IsAny <CancellationToken>()))
.ReturnsAsync(TryCatch <PartitionedQueryExecutionInfo> .FromException(
new InvalidOperationException(
exceptionMessage)));
CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters(
sqlQuerySpec: sqlQuerySpec,
initialUserContinuationToken: null,
initialFeedRange: null,
maxConcurrency: queryRequestOptions?.MaxConcurrency,
maxItemCount: queryRequestOptions?.MaxItemCount,
maxBufferedItemCount: queryRequestOptions?.MaxBufferedItemCount,
partitionKey: queryRequestOptions?.PartitionKey,
properties: queryRequestOptions?.Properties,
partitionedQueryExecutionInfo: null,
executionEnvironment: queryRequestOptions?.ExecutionEnvironment,
returnResultsInDeterministicOrder: true,
forcePassthrough: false,
testInjections: queryRequestOptions?.TestSettings);
CosmosQueryContext cosmosQueryContext = new CosmosQueryContextCore(
client: client.Object,
queryRequestOptions: queryRequestOptions,
resourceTypeEnum: ResourceType.Document,
operationType: OperationType.Query,
resourceType: typeof(QueryResponse),
resourceLink: "dbs/mockdb/colls/mockColl",
isContinuationExpected: isContinuationExpected,
allowNonValueAggregateQuery: allowNonValueAggregateQuery,
diagnosticsContext: new CosmosDiagnosticsContextCore(),
correlatedActivityId: new Guid("221FC86C-1825-4284-B10E-A6029652CCA6"));
IQueryPipelineStage pipelineStage = CosmosQueryExecutionContextFactory.Create(
documentContainer: null,
cosmosQueryContext,
inputParameters);
Assert.IsTrue(await pipelineStage.MoveNextAsync());
TryCatch <QueryPage> tryGetPage = pipelineStage.Current;
Assert.IsTrue(tryGetPage.Failed);
Assert.AreEqual(HttpStatusCode.BadRequest, (tryGetPage.InnerMostException as CosmosException).StatusCode);
Assert.IsTrue(
(tryGetPage.InnerMostException as CosmosException).ToString().Contains(exceptionMessage),
"response error message did not contain the proper substring.");
}
public EnumerableStage(IQueryPipelineStage stage)
{
this.stage = stage ?? throw new ArgumentNullException(nameof(stage));
}
