private async Task <QueryResponseCore> GetDecryptedElementResponseAsync(
QueryResponseCore queryResponseCore,
ResponseMessage message,
CancellationToken cancellationToken)
{
List <CosmosElement> documents = new List <CosmosElement>();
using (message.DiagnosticsContext.CreateScope("Decrypt"))
{
foreach (CosmosElement document in queryResponseCore.CosmosElements)
{
if (!(document is CosmosObject documentObject))
{
documents.Add(document);
continue;
}
CosmosObject decryptedDocument = await this.clientContext.EncryptionProcessor.DecryptAsync(
documentObject,
this.clientContext.ClientOptions.Encryptor,
message.DiagnosticsContext,
cancellationToken);
documents.Add(decryptedDocument);
}
}
return(QueryResponseCore.CreateSuccess(
result: documents,
requestCharge: queryResponseCore.RequestCharge,
activityId: queryResponseCore.ActivityId,
responseLengthBytes: queryResponseCore.ResponseLengthBytes,
disallowContinuationTokenMessage: queryResponseCore.DisallowContinuationTokenMessage,
continuationToken: queryResponseCore.ContinuationToken));
}
/// <summary>
/// Drains a page of results returning only distinct elements.
/// </summary>
/// <param name="maxElements">The maximum number of items to drain.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A page of distinct results.</returns>
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
List <CosmosElement> distinctResults = new List <CosmosElement>();
QueryResponseCore sourceResponse = await base.DrainAsync(maxElements, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
foreach (CosmosElement document in sourceResponse.CosmosElements)
{
if (this.distinctMap.Add(document, out UInt128 hash))
{
distinctResults.Add(document);
}
}
// For clients we write out the continuation token if it's a streaming query.
QueryResponseCore queryResponseCore;
if (this.distinctQueryType == DistinctQueryType.Ordered)
{
string updatedContinuationToken;
if (this.IsDone)
{
updatedContinuationToken = null;
}
else
{
updatedContinuationToken = new DistinctContinuationToken(
sourceToken: sourceResponse.ContinuationToken,
distinctMapToken: this.distinctMap.GetContinuationToken()).ToString();
}
queryResponseCore = QueryResponseCore.CreateSuccess(
result: distinctResults,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: null,
activityId: sourceResponse.ActivityId,
requestCharge: sourceResponse.RequestCharge,
diagnostics: sourceResponse.Diagnostics,
responseLengthBytes: sourceResponse.ResponseLengthBytes);
}
else
{
queryResponseCore = QueryResponseCore.CreateSuccess(
result: distinctResults,
continuationToken: null,
disallowContinuationTokenMessage: ClientDistinctDocumentQueryExecutionComponent.DisallowContinuationTokenMessage,
activityId: sourceResponse.ActivityId,
requestCharge: sourceResponse.RequestCharge,
diagnostics: sourceResponse.Diagnostics,
responseLengthBytes: sourceResponse.ResponseLengthBytes);
}
return(queryResponseCore);
}
/// <summary>
/// Drains a page of results returning only distinct elements.
/// </summary>
/// <param name="maxElements">The maximum number of items to drain.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A page of distinct results.</returns>
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken cancellationToken)
{
List <CosmosElement> distinctResults = new List <CosmosElement>();
QueryResponseCore sourceResponse = await base.DrainAsync(maxElements, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
foreach (CosmosElement document in sourceResponse.CosmosElements)
{
if (this.distinctMap.Add(document, out UInt128 hash))
{
distinctResults.Add(document);
}
}
return(QueryResponseCore.CreateSuccess(
result: distinctResults,
continuationToken: null,
disallowContinuationTokenMessage: ComputeDistinctDocumentQueryExecutionComponent.UseTryGetContinuationTokenMessage,
activityId: sourceResponse.ActivityId,
requestCharge: sourceResponse.RequestCharge,
responseLengthBytes: sourceResponse.ResponseLengthBytes));
}
public static (QueryResponseCore queryResponse, IList <ToDoItem> items) Create(
string itemIdPrefix,
string continuationToken,
string collectionRid,
int itemCount = 50)
{
// Use -1 to represent a split response
if (itemCount == QueryResponseMessageFactory.SPLIT)
{
return(CreateSplitResponse(collectionRid), new List <ToDoItem>().AsReadOnly());
}
IList <ToDoItem> items = ToDoItem.CreateItems(itemCount, itemIdPrefix);
MemoryStream memoryStream = (MemoryStream)cosmosSerializer.ToStream <IList <ToDoItem> >(items);
long responseLengthBytes = memoryStream.Length;
IJsonNavigator jsonNavigator = JsonNavigator.Create(memoryStream.ToArray());
IJsonNavigatorNode jsonNavigatorNode = jsonNavigator.GetRootNode();
CosmosArray cosmosArray = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);
QueryResponseCore message = QueryResponseCore.CreateSuccess(
result: cosmosArray,
continuationToken: continuationToken,
disallowContinuationTokenMessage: null,
activityId: Guid.NewGuid().ToString(),
requestCharge: 42,
queryMetricsText: null,
queryMetrics: new Dictionary <string, QueryMetrics>(),
requestStatistics: null,
responseLengthBytes: responseLengthBytes);
return(message, items);
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore results = await base.DrainAsync(maxElements, token);
if (!results.IsSuccess)
{
return(results);
}
List <CosmosElement> takedDocuments = results.CosmosElements.Take(this.takeCount).ToList();
this.takeCount -= takedDocuments.Count;
string updatedContinuationToken = null;
if (results.DisallowContinuationTokenMessage == null)
{
if (!this.TryGetContinuationToken(out updatedContinuationToken))
{
throw new InvalidOperationException($"Failed to get state for {nameof(TakeDocumentQueryExecutionComponent)}.");
}
}
return(QueryResponseCore.CreateSuccess(
result: takedDocuments,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: results.DisallowContinuationTokenMessage,
activityId: results.ActivityId,
requestCharge: results.RequestCharge,
diagnostics: results.Diagnostics,
responseLengthBytes: results.ResponseLengthBytes));
}
public static QueryResponseCore CreateQueryResponse(
IList <ToDoItem> items,
bool isOrderByQuery,
string continuationToken,
string collectionRid)
{
MemoryStream memoryStream;
string json;
if (isOrderByQuery)
{
memoryStream = SerializeForOrderByQuery(items);
using (StreamReader sr = new StreamReader(SerializeForOrderByQuery(items)))
{
json = sr.ReadToEnd();
}
}
else
{
memoryStream = (MemoryStream)MockCosmosUtil.Serializer.ToStream <IList <ToDoItem> >(items);
}
long responseLengthBytes = memoryStream.Length;
IJsonNavigator jsonNavigator = JsonNavigator.Create(memoryStream.ToArray());
IJsonNavigatorNode jsonNavigatorNode = jsonNavigator.GetRootNode();
CosmosArray cosmosArray = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);
CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();
diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
activityId: Guid.NewGuid().ToString(),
statusCode: HttpStatusCode.OK,
subStatusCode: SubStatusCodes.Unknown,
responseTimeUtc: DateTime.UtcNow,
requestCharge: 4,
errorMessage: null,
method: HttpMethod.Post,
requestUri: "http://localhost.com",
requestSessionToken: null,
responseSessionToken: null));
IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
{
new QueryPageDiagnostics(
Guid.NewGuid(),
"0",
"SomeQueryMetricText",
"SomeIndexUtilText",
diagnosticsContext)
};
QueryResponseCore message = QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: 4,
activityId: Guid.NewGuid().ToString(),
responseLengthBytes: responseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: continuationToken);
return(message);
}
private static QueryResponseCore GetCosmosElementResponse(
Guid clientQueryCorrelationId,
QueryRequestOptions requestOptions,
ResourceType resourceType,
ResponseMessage cosmosResponseMessage,
PartitionKeyRangeIdentity partitionKeyRangeIdentity,
Action <QueryPageDiagnostics> queryPageDiagnostics)
{
using (cosmosResponseMessage)
{
QueryPageDiagnostics queryPage = new QueryPageDiagnostics(
clientQueryCorrelationId: clientQueryCorrelationId,
partitionKeyRangeId: partitionKeyRangeIdentity.PartitionKeyRangeId,
queryMetricText: cosmosResponseMessage.Headers.QueryMetricsText,
indexUtilizationText: cosmosResponseMessage.Headers[HttpConstants.HttpHeaders.IndexUtilization],
diagnosticsContext: cosmosResponseMessage.DiagnosticsContext);
queryPageDiagnostics(queryPage);
if (!cosmosResponseMessage.IsSuccessStatusCode)
{
return(QueryResponseCore.CreateFailure(
statusCode: cosmosResponseMessage.StatusCode,
subStatusCodes: cosmosResponseMessage.Headers.SubStatusCode,
cosmosException: cosmosResponseMessage.CosmosException,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId));
}
if (!(cosmosResponseMessage.Content is MemoryStream memoryStream))
{
memoryStream = new MemoryStream();
cosmosResponseMessage.Content.CopyTo(memoryStream);
}
long responseLengthBytes = memoryStream.Length;
CosmosArray cosmosArray = CosmosQueryClientCore.ParseElementsFromRestStream(
memoryStream,
resourceType,
requestOptions.CosmosSerializationFormatOptions);
CosmosQueryExecutionInfo cosmosQueryExecutionInfo;
if (cosmosResponseMessage.Headers.TryGetValue(QueryExecutionInfoHeader, out string queryExecutionInfoString))
{
cosmosQueryExecutionInfo = JsonConvert.DeserializeObject <CosmosQueryExecutionInfo>(queryExecutionInfoString);
}
else
{
cosmosQueryExecutionInfo = default;
}
return(QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
responseLengthBytes: responseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: cosmosResponseMessage.Headers.ContinuationToken,
cosmosQueryExecutionInfo: cosmosQueryExecutionInfo));
}
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token);
if (!sourcePage.IsSuccess)
{
return(sourcePage);
}
// skip the documents but keep all the other headers
IReadOnlyList <CosmosElement> documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList();
int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count();
this.skipCount -= numberOfDocumentsSkipped;
return(QueryResponseCore.CreateSuccess(
result: documentsAfterSkip,
continuationToken: null,
disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead,
activityId: sourcePage.ActivityId,
requestCharge: sourcePage.RequestCharge,
diagnostics: sourcePage.Diagnostics,
responseLengthBytes: sourcePage.ResponseLengthBytes));
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token);
if (!sourcePage.IsSuccess)
{
return(sourcePage);
}
// skip the documents but keep all the other headers
IReadOnlyList <CosmosElement> documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList();
int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count() - documentsAfterSkip.Count();
this.skipCount -= numberOfDocumentsSkipped;
string updatedContinuationToken = null;
if (sourcePage.DisallowContinuationTokenMessage == null)
{
if (!this.TryGetContinuationToken(out updatedContinuationToken))
{
throw new InvalidOperationException($"Failed to get state for {nameof(SkipDocumentQueryExecutionComponent)}.");
}
}
return(QueryResponseCore.CreateSuccess(
result: documentsAfterSkip,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage,
activityId: sourcePage.ActivityId,
requestCharge: sourcePage.RequestCharge,
diagnostics: sourcePage.Diagnostics,
responseLengthBytes: sourcePage.ResponseLengthBytes));
}
public override async Task <QueryResponseCore> DrainAsync(
int maxElements,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// Note-2016-10-25-felixfan: Given what we support now, we should expect to return only 1 document.
// Note-2019-07-11-brchon: We can return empty pages until all the documents are drained,
// but then we will have to design a continuation token.
double requestCharge = 0;
long responseLengthBytes = 0;
List <QueryPageDiagnostics> diagnosticsPages = new List <QueryPageDiagnostics>();
while (!this.Source.IsDone)
{
QueryResponseCore sourceResponse = await this.Source.DrainAsync(int.MaxValue, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
requestCharge += sourceResponse.RequestCharge;
responseLengthBytes += sourceResponse.ResponseLengthBytes;
if (sourceResponse.Diagnostics != null)
{
diagnosticsPages.AddRange(sourceResponse.Diagnostics);
}
foreach (CosmosElement element in sourceResponse.CosmosElements)
{
RewrittenAggregateProjections rewrittenAggregateProjections = new RewrittenAggregateProjections(
this.isValueAggregateQuery,
element);
this.singleGroupAggregator.AddValues(rewrittenAggregateProjections.Payload);
}
}
List <CosmosElement> finalResult = new List <CosmosElement>();
CosmosElement aggregationResult = this.singleGroupAggregator.GetResult();
if (aggregationResult != null)
{
finalResult.Add(aggregationResult);
}
return(QueryResponseCore.CreateSuccess(
result: finalResult,
continuationToken: null,
activityId: null,
disallowContinuationTokenMessage: null,
requestCharge: requestCharge,
diagnostics: diagnosticsPages,
responseLengthBytes: responseLengthBytes));
}
private QueryResponseCore GetCosmosElementResponse(
QueryRequestOptions requestOptions,
ResourceType resourceType,
ResponseMessage cosmosResponseMessage,
PartitionKeyRangeIdentity partitionKeyRangeIdentity,
SchedulingStopwatch schedulingStopwatch)
{
using (cosmosResponseMessage)
{
QueryPageDiagnostics diagnostics = new QueryPageDiagnostics(
partitionKeyRangeId: partitionKeyRangeIdentity.PartitionKeyRangeId,
queryMetricText: cosmosResponseMessage.Headers.QueryMetricsText,
indexUtilizationText: cosmosResponseMessage.Headers[HttpConstants.HttpHeaders.IndexUtilization],
requestDiagnostics: cosmosResponseMessage.Diagnostics,
schedulingStopwatch: schedulingStopwatch);
IReadOnlyCollection <QueryPageDiagnostics> pageDiagnostics = new List <QueryPageDiagnostics>()
{
diagnostics
};
if (!cosmosResponseMessage.IsSuccessStatusCode)
{
return(QueryResponseCore.CreateFailure(
statusCode: cosmosResponseMessage.StatusCode,
subStatusCodes: cosmosResponseMessage.Headers.SubStatusCode,
errorMessage: cosmosResponseMessage.ErrorMessage,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
diagnostics: pageDiagnostics));
}
MemoryStream memoryStream = cosmosResponseMessage.Content as MemoryStream;
if (memoryStream == null)
{
memoryStream = new MemoryStream();
cosmosResponseMessage.Content.CopyTo(memoryStream);
}
long responseLengthBytes = memoryStream.Length;
CosmosArray cosmosArray = CosmosElementSerializer.ToCosmosElements(
memoryStream,
resourceType,
requestOptions.CosmosSerializationFormatOptions);
int itemCount = cosmosArray.Count;
return(QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
diagnostics: pageDiagnostics,
responseLengthBytes: cosmosResponseMessage.Headers.ContentLengthAsLong,
disallowContinuationTokenMessage: null,
continuationToken: cosmosResponseMessage.Headers.ContinuationToken));
}
}
public static (QueryResponseCore queryResponse, IList <ToDoItem> items) Create(
string itemIdPrefix,
string continuationToken,
string collectionRid,
int itemCount = 50)
{
// Use -1 to represent a split response
if (itemCount == QueryResponseMessageFactory.SPLIT)
{
return(CreateSplitResponse(collectionRid), new List <ToDoItem>().AsReadOnly());
}
IList <ToDoItem> items = ToDoItem.CreateItems(itemCount, itemIdPrefix);
MemoryStream memoryStream = (MemoryStream)MockCosmosUtil.Serializer.ToStream <IList <ToDoItem> >(items);
long responseLengthBytes = memoryStream.Length;
IJsonNavigator jsonNavigator = JsonNavigator.Create(memoryStream.ToArray());
IJsonNavigatorNode jsonNavigatorNode = jsonNavigator.GetRootNode();
CosmosArray cosmosArray = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);
double requestCharge = 42;
string activityId = Guid.NewGuid().ToString();
CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();
diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
activityId: Guid.NewGuid().ToString(),
statusCode: HttpStatusCode.OK,
subStatusCode: SubStatusCodes.Unknown,
responseTimeUtc: DateTime.UtcNow,
requestCharge: requestCharge,
errorMessage: null,
method: HttpMethod.Post,
requestUri: "http://localhost.com",
requestSessionToken: null,
responseSessionToken: null));
IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
{
new QueryPageDiagnostics(
Guid.NewGuid(),
"0",
"SomeQueryMetricText",
"SomeIndexUtilText",
diagnosticsContext)
};
QueryResponseCore message = QueryResponseCore.CreateSuccess(
result: cosmosArray,
continuationToken: continuationToken,
disallowContinuationTokenMessage: null,
activityId: activityId,
requestCharge: requestCharge,
responseLengthBytes: responseLengthBytes);
return(message, items);
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore results = await base.DrainAsync(maxElements, token);
if (!results.IsSuccess)
{
return(results);
}
List <CosmosElement> takedDocuments = results.CosmosElements.Take(this.takeCount).ToList();
this.takeCount -= takedDocuments.Count;
string updatedContinuationToken = null;
if (results.DisallowContinuationTokenMessage == null)
{
if (!this.IsDone)
{
string sourceContinuation = results.ContinuationToken;
TakeContinuationToken takeContinuationToken;
switch (this.takeEnum)
{
case TakeEnum.Limit:
takeContinuationToken = new LimitContinuationToken(
this.takeCount,
sourceContinuation);
break;
case TakeEnum.Top:
takeContinuationToken = new TopContinuationToken(
this.takeCount,
sourceContinuation);
break;
default:
throw new ArgumentException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}");
}
updatedContinuationToken = takeContinuationToken.ToString();
}
}
return(QueryResponseCore.CreateSuccess(
result: takedDocuments,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: results.DisallowContinuationTokenMessage,
activityId: results.ActivityId,
requestCharge: results.RequestCharge,
queryMetricsText: results.QueryMetricsText,
queryMetrics: results.QueryMetrics,
requestStatistics: results.RequestStatistics,
responseLengthBytes: results.ResponseLengthBytes));
}
private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse)
{
// We need to give empty pages until the results are fully drained.
QueryResponseCore response = QueryResponseCore.CreateSuccess(
result: EmptyResults,
requestCharge: sourceResponse.RequestCharge,
activityId: sourceResponse.ActivityId,
responseLengthBytes: sourceResponse.ResponseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: sourceResponse.ContinuationToken);
return(response);
}
private QueryResponseCore GetCosmosElementResponse(
Guid clientQueryCorrelationId,
QueryRequestOptions requestOptions,
ResourceType resourceType,
ResponseMessage cosmosResponseMessage,
PartitionKeyRangeIdentity partitionKeyRangeIdentity,
Action <QueryPageDiagnostics> queryPageDiagnostics)
{
using (cosmosResponseMessage)
{
QueryPageDiagnostics queryPage = new QueryPageDiagnostics(
clientQueryCorrelationId: clientQueryCorrelationId,
partitionKeyRangeId: partitionKeyRangeIdentity.PartitionKeyRangeId,
queryMetricText: cosmosResponseMessage.Headers.QueryMetricsText,
indexUtilizationText: cosmosResponseMessage.Headers[HttpConstants.HttpHeaders.IndexUtilization],
diagnosticsContext: cosmosResponseMessage.DiagnosticsContext);
queryPageDiagnostics(queryPage);
if (!cosmosResponseMessage.IsSuccessStatusCode)
{
return(QueryResponseCore.CreateFailure(
statusCode: cosmosResponseMessage.StatusCode,
subStatusCodes: cosmosResponseMessage.Headers.SubStatusCode,
cosmosException: cosmosResponseMessage.CosmosException,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId));
}
if (!(cosmosResponseMessage.Content is MemoryStream memoryStream))
{
memoryStream = new MemoryStream();
cosmosResponseMessage.Content.CopyTo(memoryStream);
}
long responseLengthBytes = memoryStream.Length;
CosmosArray cosmosArray = CosmosElementSerializer.ToCosmosElements(
memoryStream,
resourceType,
requestOptions.CosmosSerializationFormatOptions);
int itemCount = cosmosArray.Count;
return(QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
responseLengthBytes: responseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: cosmosResponseMessage.Headers.ContinuationToken));
}
}
public override async Task <QueryResponseCore> DrainAsync(
int maxElements,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// Draining GROUP BY is broken down into two stages:
QueryResponseCore response;
if (!this.Source.IsDone)
{
// Stage 1:
// Drain the groupings fully from all continuation and all partitions
QueryResponseCore sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
this.AggregateGroupings(sourceResponse.CosmosElements);
// We need to give empty pages until the results are fully drained.
response = QueryResponseCore.CreateSuccess(
result: EmptyResults,
continuationToken: null,
disallowContinuationTokenMessage: UseTryGetContinuationTokenInstead,
activityId: sourceResponse.ActivityId,
requestCharge: sourceResponse.RequestCharge,
diagnostics: sourceResponse.Diagnostics,
responseLengthBytes: sourceResponse.ResponseLengthBytes);
}
else
{
// Stage 2:
// Emit the results from the grouping table page by page
IReadOnlyList <CosmosElement> results = this.groupingTable.Drain(maxElements);
response = QueryResponseCore.CreateSuccess(
result: results,
continuationToken: null,
disallowContinuationTokenMessage: UseTryGetContinuationTokenInstead,
activityId: null,
requestCharge: 0,
diagnostics: QueryResponseCore.EmptyDiagnostics,
responseLengthBytes: 0);
}
return(response);
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token);
if (!sourcePage.IsSuccess)
{
return(sourcePage);
}
List <CosmosElement> takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList();
this.takeCount -= takedDocuments.Count;
string updatedContinuationToken;
if (!this.IsDone && (sourcePage.DisallowContinuationTokenMessage == null))
{
switch (this.takeEnum)
{
case TakeEnum.Limit:
updatedContinuationToken = new LimitContinuationToken(
limit: this.takeCount,
sourceToken: sourcePage.ContinuationToken).ToString();
break;
case TakeEnum.Top:
updatedContinuationToken = new TopContinuationToken(
top: this.takeCount,
sourceToken: sourcePage.ContinuationToken).ToString();
break;
default:
throw new ArgumentOutOfRangeException($"Unknown {nameof(TakeEnum)}: {this.takeEnum}.");
}
}
else
{
updatedContinuationToken = null;
}
return(QueryResponseCore.CreateSuccess(
result: takedDocuments,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage,
activityId: sourcePage.ActivityId,
requestCharge: sourcePage.RequestCharge,
diagnostics: sourcePage.Diagnostics,
responseLengthBytes: sourcePage.ResponseLengthBytes));
}
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()));
}
public override async Task <QueryResponseCore> DrainAsync(
int maxElements,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// Draining GROUP BY is broken down into two stages:
double requestCharge = 0.0;
long responseLengthBytes = 0;
List <QueryPageDiagnostics> queryPageDiagnostics = new List <QueryPageDiagnostics>();
while (!this.Source.IsDone)
{
// Stage 1:
// Drain the groupings fully from all continuation and all partitions
QueryResponseCore sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
requestCharge += sourceResponse.RequestCharge;
responseLengthBytes += sourceResponse.ResponseLengthBytes;
if (sourceResponse.Diagnostics != null)
{
queryPageDiagnostics.AddRange(sourceResponse.Diagnostics);
}
this.AggregateGroupings(sourceResponse.CosmosElements);
}
// Stage 2:
// Emit the results from the grouping table page by page
IReadOnlyList <CosmosElement> results = this.groupingTable.Drain(maxElements);
QueryResponseCore response = QueryResponseCore.CreateSuccess(
result: results,
continuationToken: null,
disallowContinuationTokenMessage: ClientGroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy,
activityId: null,
requestCharge: requestCharge,
diagnostics: queryPageDiagnostics,
responseLengthBytes: responseLengthBytes);
return(response);
}
/// <summary>
/// Drains a page of results returning only distinct elements.
/// </summary>
/// <param name="maxElements">The maximum number of items to drain.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A page of distinct results.</returns>
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken cancellationToken)
{
List <CosmosElement> distinctResults = new List <CosmosElement>();
QueryResponseCore cosmosQueryResponse = await base.DrainAsync(maxElements, cancellationToken);
if (!cosmosQueryResponse.IsSuccess)
{
return(cosmosQueryResponse);
}
foreach (CosmosElement document in cosmosQueryResponse.CosmosElements)
{
if (this.distinctMap.Add(document, out this.lastHash))
{
distinctResults.Add(document);
}
}
string updatedContinuationToken;
if (!this.IsDone)
{
updatedContinuationToken = new DistinctContinuationToken(
this.lastHash,
cosmosQueryResponse.ContinuationToken).ToString();
}
else
{
this.Source.Stop();
updatedContinuationToken = null;
}
string disallowContinuationTokenMessage = this.distinctQueryType == DistinctQueryType.Ordered ? null : RMResources.UnorderedDistinctQueryContinuationToken;
return(QueryResponseCore.CreateSuccess(
result: distinctResults,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: disallowContinuationTokenMessage,
activityId: cosmosQueryResponse.ActivityId,
requestCharge: cosmosQueryResponse.RequestCharge,
queryMetricsText: cosmosQueryResponse.QueryMetricsText,
queryMetrics: cosmosQueryResponse.QueryMetrics,
requestStatistics: cosmosQueryResponse.RequestStatistics,
responseLengthBytes: cosmosQueryResponse.ResponseLengthBytes));
}
private QueryResponseCore GetCosmosElementResponse(
QueryRequestOptions requestOptions,
ResourceType resourceType,
ResponseMessage cosmosResponseMessage)
{
using (cosmosResponseMessage)
{
if (!cosmosResponseMessage.IsSuccessStatusCode)
{
return(QueryResponseCore.CreateFailure(
statusCode: cosmosResponseMessage.StatusCode,
subStatusCodes: cosmosResponseMessage.Headers.SubStatusCode,
errorMessage: cosmosResponseMessage.ErrorMessage,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
queryMetricsText: cosmosResponseMessage.Headers.QueryMetricsText,
queryMetrics: null));
}
MemoryStream memoryStream = cosmosResponseMessage.Content as MemoryStream;
if (memoryStream == null)
{
memoryStream = new MemoryStream();
cosmosResponseMessage.Content.CopyTo(memoryStream);
}
long responseLengthBytes = memoryStream.Length;
CosmosArray cosmosArray = CosmosElementSerializer.ToCosmosElements(
memoryStream,
resourceType,
requestOptions.CosmosSerializationFormatOptions);
int itemCount = cosmosArray.Count;
return(QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: cosmosResponseMessage.Headers.RequestCharge,
activityId: cosmosResponseMessage.Headers.ActivityId,
queryMetricsText: cosmosResponseMessage.Headers.QueryMetricsText,
queryMetrics: null,
requestStatistics: null,
responseLengthBytes: cosmosResponseMessage.Headers.ContentLengthAsLong,
disallowContinuationTokenMessage: null,
continuationToken: cosmosResponseMessage.Headers.ContinuationToken));
}
}
/// <summary>
/// Gets the next page of results from this context.
/// </summary>
/// <param name="token">The cancellation token.</param>
/// <returns>A task to await on that in turn returns a DoucmentFeedResponse of results.</returns>
public override async Task <QueryResponseCore> ExecuteNextAsync(CancellationToken token)
{
try
{
QueryResponseCore queryResponse = await this.component.DrainAsync(this.actualPageSize, token);
if (!queryResponse.IsSuccess)
{
this.component.Stop();
return(queryResponse);
}
string updatedContinuationToken;
if (queryResponse.DisallowContinuationTokenMessage == null)
{
if (queryResponse.ContinuationToken != null)
{
updatedContinuationToken = new PipelineContinuationTokenV0(queryResponse.ContinuationToken).ToString();
}
else
{
updatedContinuationToken = null;
}
}
else
{
updatedContinuationToken = null;
}
return(QueryResponseCore.CreateSuccess(
result: queryResponse.CosmosElements,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: queryResponse.DisallowContinuationTokenMessage,
activityId: queryResponse.ActivityId,
requestCharge: queryResponse.RequestCharge,
diagnostics: queryResponse.Diagnostics,
responseLengthBytes: queryResponse.ResponseLengthBytes));
}
catch (Exception)
{
this.component.Stop();
throw;
}
}
private QueryResponseCore GetEmptyPage(QueryResponseCore sourceResponse)
{
if (!this.TryGetContinuationToken(out string updatedContinuationToken))
{
throw new InvalidOperationException("Failed to get source continuation token.");
}
// We need to give empty pages until the results are fully drained.
QueryResponseCore response = QueryResponseCore.CreateSuccess(
result: EmptyResults,
requestCharge: sourceResponse.RequestCharge,
activityId: sourceResponse.ActivityId,
responseLengthBytes: sourceResponse.ResponseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: updatedContinuationToken,
diagnostics: sourceResponse.Diagnostics);
return(response);
}
private QueryResponseCore GetFinalResponse()
{
List <CosmosElement> finalResult = new List <CosmosElement>();
CosmosElement aggregationResult = this.singleGroupAggregator.GetResult();
if (aggregationResult != null)
{
finalResult.Add(aggregationResult);
}
QueryResponseCore response = QueryResponseCore.CreateSuccess(
result: finalResult,
requestCharge: 0,
activityId: null,
responseLengthBytes: 0,
disallowContinuationTokenMessage: null,
continuationToken: null);
return(response);
}
public static QueryResponseCore CreateQueryResponse(
IList <ToDoItem> items,
bool isOrderByQuery,
string continuationToken,
string collectionRid)
{
MemoryStream memoryStream;
string json;
if (isOrderByQuery)
{
memoryStream = SerializeForOrderByQuery(items);
using (StreamReader sr = new StreamReader(SerializeForOrderByQuery(items)))
{
json = sr.ReadToEnd();
}
}
else
{
memoryStream = (MemoryStream)cosmosSerializer.ToStream <IList <ToDoItem> >(items);
}
long responseLengthBytes = memoryStream.Length;
IJsonNavigator jsonNavigator = JsonNavigator.Create(memoryStream.ToArray());
IJsonNavigatorNode jsonNavigatorNode = jsonNavigator.GetRootNode();
CosmosArray cosmosArray = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);
QueryResponseCore message = QueryResponseCore.CreateSuccess(
result: cosmosArray,
requestCharge: 4,
activityId: Guid.NewGuid().ToString(),
queryMetricsText: null,
queryMetrics: null,
requestStatistics: null,
responseLengthBytes: responseLengthBytes,
disallowContinuationTokenMessage: null,
continuationToken: continuationToken);
return(message);
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token);
if (!sourcePage.IsSuccess)
{
return(sourcePage);
}
List <CosmosElement> takedDocuments = sourcePage.CosmosElements.Take(this.takeCount).ToList();
this.takeCount -= takedDocuments.Count;
return(QueryResponseCore.CreateSuccess(
result: takedDocuments,
continuationToken: null,
disallowContinuationTokenMessage: DocumentQueryExecutionComponentBase.UseCosmosElementContinuationTokenInstead,
activityId: sourcePage.ActivityId,
requestCharge: sourcePage.RequestCharge,
responseLengthBytes: sourcePage.ResponseLengthBytes));
}
public override async Task <QueryResponseCore> DrainAsync(int maxElements, CancellationToken token)
{
token.ThrowIfCancellationRequested();
QueryResponseCore sourcePage = await base.DrainAsync(maxElements, token);
if (!sourcePage.IsSuccess)
{
return(sourcePage);
}
// skip the documents but keep all the other headers
List <CosmosElement> documentsAfterSkip = sourcePage.CosmosElements.Skip(this.skipCount).ToList();
int numberOfDocumentsSkipped = sourcePage.CosmosElements.Count - documentsAfterSkip.Count;
this.skipCount -= numberOfDocumentsSkipped;
string updatedContinuationToken = null;
if (sourcePage.DisallowContinuationTokenMessage == null)
{
if (!this.IsDone)
{
updatedContinuationToken = new OffsetContinuationToken(
this.skipCount,
sourcePage.ContinuationToken).ToString();
}
}
return(QueryResponseCore.CreateSuccess(
result: documentsAfterSkip,
continuationToken: updatedContinuationToken,
disallowContinuationTokenMessage: sourcePage.DisallowContinuationTokenMessage,
activityId: sourcePage.ActivityId,
requestCharge: sourcePage.RequestCharge,
queryMetricsText: sourcePage.QueryMetricsText,
queryMetrics: sourcePage.QueryMetrics,
requestStatistics: sourcePage.RequestStatistics,
responseLengthBytes: sourcePage.ResponseLengthBytes));
}
/// <summary>
/// Drains a page of documents from this context.
/// </summary>
/// <param name="maxElements">The maximum number of elements.</param>
/// <param name="cancellationToken">The cancellation token.</param>
/// <returns>A task that when awaited on return a page of documents.</returns>
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 always drain from the partition, which has the highest priority item first
//// 2) If multiple partitions have the same priority item then we drain from the left most first
//// otherwise we would need to keep track of how many of each item we drained from each partition
//// (just like parallel queries).
//// Visually that look the following case where we have three partitions that are numbered and store letters.
//// For teaching purposes I have made each item a tuple of the following form:
//// <item stored in partition, partition number>
//// So that duplicates across partitions are distinct, but duplicates within partitions are indistinguishable.
//// |-------| |-------| |-------|
//// | <a,1> | | <a,2> | | <a,3> |
//// | <a,1> | | <b,2> | | <c,3> |
//// | <a,1> | | <b,2> | | <c,3> |
//// | <d,1> | | <c,2> | | <c,3> |
//// | <d,1> | | <e,2> | | <f,3> |
//// | <e,1> | | <h,2> | | <j,3> |
//// | <f,1> | | <i,2> | | <k,3> |
//// |-------| |-------| |-------|
//// Now the correct drain order in this case is:
//// <a,1>,<a,1>,<a,1>,<a,2>,<a,3>,<b,2>,<b,2>,<c,2>,<c,3>,<c,3>,<c,3>,
//// <d,1>,<d,1>,<e,1>,<e,2>,<f,1>,<f,3>,<h,2>,<i,2>,<j,3>,<k,3>
//// In more mathematical terms
//// 1) <x, y> always comes before <z, y> where x < z
//// 2) <i, j> always come before <i, k> where j < k
List <CosmosElement> results = new List <CosmosElement>();
while (results.Count < maxElements)
{
// 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.
ItemProducerTree currentItemProducerTree = this.PopCurrentItemProducerTree();
try
{
if (!currentItemProducerTree.HasMoreResults)
{
// This means there are no more items to drain
break;
}
OrderByQueryResult orderByQueryResult = new OrderByQueryResult(currentItemProducerTree.Current);
// Only add the payload, since other stuff is garbage from the caller's perspective.
results.Add(orderByQueryResult.Payload);
// If we are at the beginning of the page and seeing an rid from the previous page we should increment the skip count
// due to the fact that JOINs can make a document appear multiple times and across continuations, so we don't want to
// surface this more than needed. More information can be found in the continuation token docs.
if (this.ShouldIncrementSkipCount(currentItemProducerTree.CurrentItemProducerTree.Root))
{
++this.skipCount;
}
else
{
this.skipCount = 0;
}
this.previousRid = orderByQueryResult.Rid;
this.previousOrderByItems = orderByQueryResult.OrderByItems;
if (!currentItemProducerTree.TryMoveNextDocumentWithinPage())
{
while (true)
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await currentItemProducerTree.TryMoveNextPageAsync(cancellationToken);
if (!movedToNextPage)
{
if (failureResponse.HasValue)
{
// TODO: We can buffer this failure so that the user can still get the pages we already got.
return(failureResponse.Value);
}
break;
}
if (currentItemProducerTree.IsAtBeginningOfPage)
{
break;
}
if (currentItemProducerTree.TryMoveNextDocumentWithinPage())
{
break;
}
}
}
}
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()));
}
public async Task TestItemProducerTreeWithFailure()
{
int callBackCount = 0;
Mock <CosmosQueryContext> mockQueryContext = new Mock <CosmosQueryContext>();
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("Select * from t");
PartitionKeyRange partitionKeyRange = new PartitionKeyRange {
Id = "0", MinInclusive = "A", MaxExclusive = "B"
};
void produceAsyncCompleteCallback(
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
IReadOnlyCollection <QueryPageDiagnostics> queryPageDiagnostics,
long responseLengthBytes,
CancellationToken token)
{
callBackCount++;
}
Mock <IComparer <ItemProducerTree> > comparer = new Mock <IComparer <ItemProducerTree> >();
Mock <IEqualityComparer <CosmosElement> > cosmosElementComparer = new Mock <IEqualityComparer <CosmosElement> >();
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
IReadOnlyList <CosmosElement> cosmosElements = new List <CosmosElement>()
{
new Mock <CosmosElement>(CosmosElementType.Object).Object
};
QueryPageDiagnostics diagnostics = new QueryPageDiagnostics(
partitionKeyRangeId: "0",
queryMetricText: "SomeRandomQueryMetricText",
indexUtilizationText: null,
requestDiagnostics: new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.OK,
subStatusCode: SubStatusCodes.Unknown,
requestCharge: 42,
errorMessage: null,
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
requestSessionToken: null,
responseSessionToken: null,
clientSideRequestStatistics: null),
schedulingStopwatch: new SchedulingStopwatch());
IReadOnlyCollection <QueryPageDiagnostics> pageDiagnostics = new List <QueryPageDiagnostics>()
{
diagnostics
};
mockQueryContext.Setup(x => x.ContainerResourceId).Returns("MockCollectionRid");
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateSuccess(
result: cosmosElements,
requestCharge: 42,
activityId: "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE",
diagnostics: pageDiagnostics,
responseLengthBytes: 500,
disallowContinuationTokenMessage: null,
continuationToken: "TestToken")));
ItemProducerTree itemProducerTree = new ItemProducerTree(
queryContext: mockQueryContext.Object,
querySpecForInit: sqlQuerySpec,
partitionKeyRange: partitionKeyRange,
produceAsyncCompleteCallback: produceAsyncCompleteCallback,
itemProducerTreeComparer: comparer.Object,
equalityComparer: cosmosElementComparer.Object,
testSettings: new TestInjections(simulate429s: false, simulateEmptyPages: false),
deferFirstPage: false,
collectionRid: "collectionRid",
initialContinuationToken: null,
initialPageSize: 50);
// Buffer to success responses
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
diagnostics = new QueryPageDiagnostics(
partitionKeyRangeId: "0",
queryMetricText: null,
indexUtilizationText: null,
requestDiagnostics: new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.InternalServerError,
subStatusCode: SubStatusCodes.Unknown,
requestCharge: 10.2,
errorMessage: "Error message",
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
requestSessionToken: null,
responseSessionToken: null,
clientSideRequestStatistics: null),
schedulingStopwatch: new SchedulingStopwatch());
pageDiagnostics = new List <QueryPageDiagnostics>()
{
diagnostics
};
// Buffer a failure
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateFailure(
statusCode: HttpStatusCode.InternalServerError,
subStatusCodes: null,
errorMessage: "Error message",
requestCharge: 10.2,
activityId: Guid.NewGuid().ToString(),
diagnostics: pageDiagnostics)));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// First item should be a success
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);
Assert.IsTrue(movedToNextPage);
Assert.IsNull(failureResponse);
Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage());
Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage());
Assert.IsTrue(itemProducerTree.HasMoreResults);
}
// Second item should be a success
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);
Assert.IsTrue(movedToNextPage);
Assert.IsNull(failureResponse);
Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage());
Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage());
Assert.IsTrue(itemProducerTree.HasMoreResults);
}
// Third item should be a failure
{
(bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);
Assert.IsFalse(movedToNextPage);
Assert.IsNotNull(failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
}
// Try to buffer after failure. It should return the previous cached failure and not try to buffer again.
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).
Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again"));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
Assert.IsFalse(itemProducerTree.HasMoreResults);
}
public async Task TestItemProducerTreeWithFailure()
{
int callBackCount = 0;
Mock <CosmosQueryContext> mockQueryContext = new Mock <CosmosQueryContext>();
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("Select * from t");
PartitionKeyRange partitionKeyRange = new PartitionKeyRange {
Id = "0", MinInclusive = "A", MaxExclusive = "B"
};
ItemProducerTree.ProduceAsyncCompleteDelegate produceAsyncCompleteCallback = (
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
QueryMetrics queryMetrics,
long responseLengthBytes,
CancellationToken token) =>
{ callBackCount++; };
Mock <IComparer <ItemProducerTree> > comparer = new Mock <IComparer <ItemProducerTree> >();
Mock <IEqualityComparer <CosmosElement> > cosmosElementComparer = new Mock <IEqualityComparer <CosmosElement> >();
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
IReadOnlyList <CosmosElement> cosmosElements = new List <CosmosElement>()
{
new Mock <CosmosElement>(CosmosElementType.Object).Object
};
mockQueryContext.Setup(x => x.ContainerResourceId).Returns("MockCollectionRid");
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateSuccess(
result: cosmosElements,
requestCharge: 42,
activityId: "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE",
queryMetrics: null,
queryMetricsText: null,
requestStatistics: null,
responseLengthBytes: 500,
disallowContinuationTokenMessage: null,
continuationToken: "TestToken")));
ItemProducerTree itemProducerTree = new ItemProducerTree(
queryContext: mockQueryContext.Object,
querySpecForInit: sqlQuerySpec,
partitionKeyRange: partitionKeyRange,
produceAsyncCompleteCallback: produceAsyncCompleteCallback,
itemProducerTreeComparer: comparer.Object,
equalityComparer: cosmosElementComparer.Object,
deferFirstPage: false,
collectionRid: "collectionRid",
initialContinuationToken: null,
initialPageSize: 50);
// Buffer to success responses
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// Buffer a failure
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateFailure(
statusCode: HttpStatusCode.InternalServerError,
subStatusCodes: null,
errorMessage: "Error message",
requestCharge: 10.2,
activityId: Guid.NewGuid().ToString(),
queryMetricsText: null,
queryMetrics: null)));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// First item should be a success
(bool successfullyMovedNext, QueryResponseCore? failureResponse)result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Second item should be a success
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Third item should be a failure
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
// Try to buffer after failure. It should return the previous cached failure and not try to buffer again.
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
cancellationTokenSource.Token)).
Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again"));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
}
