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);
}
public override async Task <QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// If the cache is stale the entire execute context has incorrect values and should be recreated.
// This should only be done for the first execution.
// If results have already been pulled,
// then an error should be returned to the user,
// since it's not possible to combine query results from multiple containers.
QueryResponseCore queryResponse = await this.currentCosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
if (
(queryResponse.StatusCode == System.Net.HttpStatusCode.Gone) &&
(queryResponse.SubStatusCode == Documents.SubStatusCodes.NameCacheIsStale) &&
!this.alreadyRetried)
{
await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync(
this.cosmosQueryContext.ResourceLink.OriginalString,
cancellationToken);
this.alreadyRetried = true;
this.currentCosmosQueryExecutionContext.Dispose();
this.currentCosmosQueryExecutionContext = this.cosmosQueryExecutionContextFactory();
return(await this.ExecuteNextAsync(cancellationToken));
}
return(queryResponse);
}
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);
}
private (Func <string, Task <TryCatch <IDocumentQueryExecutionComponent> > >, QueryResponseCore) SetupBaseContextToVerifyFailureScenario()
{
CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create();
diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.Unauthorized,
subStatusCode: SubStatusCodes.PartitionKeyMismatch,
requestCharge: 4,
errorMessage: null,
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
requestSessionToken: null,
responseSessionToken: null,
clientSideRequestStatistics: null));
IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
{
new QueryPageDiagnostics(
"0",
"SomeQueryMetricText",
"SomeIndexUtilText",
diagnosticsContext,
new SchedulingStopwatch())
};
QueryResponseCore failure = QueryResponseCore.CreateFailure(
System.Net.HttpStatusCode.Unauthorized,
SubStatusCodes.PartitionKeyMismatch,
new CosmosException(
statusCodes: HttpStatusCode.Unauthorized,
message: "Random error message",
subStatusCode: default,
/// <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));
}
private (Func <string, Task <IDocumentQueryExecutionComponent> >, QueryResponseCore) SetupBaseContextToVerifyFailureScenario()
{
IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
{
new QueryPageDiagnostics("0",
"SomeQueryMetricText",
"SomeIndexUtilText",
new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.Unauthorized,
subStatusCode: SubStatusCodes.PartitionKeyMismatch,
requestCharge: 4,
errorMessage: null,
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
clientSideRequestStatistics: null),
new SchedulingStopwatch())
};
QueryResponseCore failure = QueryResponseCore.CreateFailure(
System.Net.HttpStatusCode.Unauthorized,
SubStatusCodes.PartitionKeyMismatch,
"Random error message",
42.89,
"TestActivityId",
diagnostics);
Mock <IDocumentQueryExecutionComponent> baseContext = new Mock <IDocumentQueryExecutionComponent>();
baseContext.Setup(x => x.DrainAsync(It.IsAny <int>(), It.IsAny <CancellationToken>())).Returns(Task.FromResult <QueryResponseCore>(failure));
Func <string, Task <IDocumentQueryExecutionComponent> > callBack = x => Task.FromResult <IDocumentQueryExecutionComponent>(baseContext.Object);
return(callBack, failure);
}
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));
}
private static QueryResponseCore CreateFromExceptionWithStackTrace(ExceptionWithStackTraceException exceptionWithStackTrace)
{
// Use the original stack trace from the inner exception.
if (exceptionWithStackTrace.InnerException is DocumentClientException ||
exceptionWithStackTrace.InnerException is CosmosException)
{
return(QueryResponseFactory.CreateFromException(exceptionWithStackTrace.InnerException));
}
QueryResponseCore queryResponseCore = QueryResponseFactory.CreateFromException(exceptionWithStackTrace.InnerException);
CosmosException cosmosException = queryResponseCore.CosmosException;
queryResponseCore = QueryResponseCore.CreateFailure(
statusCode: queryResponseCore.StatusCode,
subStatusCodes: queryResponseCore.SubStatusCode,
cosmosException: CosmosExceptionFactory.Create(
statusCode: cosmosException.StatusCode,
message: cosmosException.Message,
stackTrace: exceptionWithStackTrace.StackTrace,
headers: cosmosException.Headers,
trace: cosmosException.Trace,
error: cosmosException.Error,
innerException: cosmosException.InnerException),
requestCharge: queryResponseCore.RequestCharge,
activityId: queryResponseCore.ActivityId);
return(queryResponseCore);
}
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 Task <QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken)
{
QueryResponseCore queryResponse = QueryResponseFactory.CreateFromException(this.exception);
this.returnedErrorResponse = true;
return(Task.FromResult(queryResponse));
}
public static QueryResponseCore CreateFailureResponse(
HttpStatusCode httpStatusCode,
SubStatusCodes subStatusCodes,
string errorMessage)
{
IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
{
new QueryPageDiagnostics("0",
"SomeQueryMetricText",
"SomeIndexUtilText",
new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.Gone,
subStatusCode: SubStatusCodes.PartitionKeyRangeGone,
requestCharge: 10.4,
errorMessage: null,
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
requestSessionToken: null,
responseSessionToken: null,
clientSideRequestStatistics: null),
new SchedulingStopwatch())
};
QueryResponseCore splitResponse = QueryResponseCore.CreateFailure(
statusCode: httpStatusCode,
subStatusCodes: subStatusCodes,
errorMessage: errorMessage,
requestCharge: 10.4,
activityId: Guid.NewGuid().ToString(),
diagnostics: diagnostics);
return(splitResponse);
}
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 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 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);
}
public void VerifyCosmosQueryResponseStream()
{
string contianerRid = "mockContainerRid";
(QueryResponseCore response, IList <ToDoItem> items)factoryResponse = QueryResponseMessageFactory.Create(
itemIdPrefix: $"TestPage",
continuationToken: "SomeContinuationToken",
collectionRid: contianerRid,
itemCount: 100);
QueryResponseCore responseCore = factoryResponse.response;
QueryResponse queryResponse = QueryResponse.CreateSuccess(
result: responseCore.CosmosElements,
count: responseCore.CosmosElements.Count,
responseLengthBytes: responseCore.ResponseLengthBytes,
queryMetrics: responseCore.QueryMetrics,
responseHeaders: new CosmosQueryResponseMessageHeaders(
responseCore.ContinuationToken,
responseCore.DisallowContinuationTokenMessage,
ResourceType.Document,
contianerRid)
{
RequestCharge = responseCore.RequestCharge,
ActivityId = responseCore.ActivityId
});
using (Stream stream = queryResponse.Content)
{
using (Stream innerStream = queryResponse.Content)
{
Assert.IsTrue(object.ReferenceEquals(stream, innerStream), "Content should return the same stream");
}
}
}
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));
}
public async Task TestCosmosQueryExecutionComponentOnFailure()
{
(IList <IDocumentQueryExecutionComponent> components, QueryResponseCore response)setupContext = await this.GetAllExecutionComponents();
foreach (DocumentQueryExecutionComponentBase component in setupContext.components)
{
QueryResponseCore response = await component.DrainAsync(1, default(CancellationToken));
Assert.AreEqual(setupContext.response, response);
}
}
private static QueryResponseCore CreateFromCosmosException(CosmosException cosmosException)
{
QueryResponseCore queryResponseCore = QueryResponseCore.CreateFailure(
statusCode: cosmosException.StatusCode,
subStatusCodes: (Microsoft.Azure.Documents.SubStatusCodes)cosmosException.SubStatusCode,
cosmosException: cosmosException,
requestCharge: 0,
activityId: cosmosException.ActivityId);
return(queryResponseCore);
}
internal override async Task <QueryResponseCore> ExecuteItemQueryAsync(
Uri resourceUri,
ResourceType resourceType,
OperationType operationType,
QueryRequestOptions requestOptions,
Action <QueryPageDiagnostics> queryPageDiagnostics,
SqlQuerySpec sqlQuerySpec,
string continuationToken,
PartitionKeyRangeIdentity partitionKeyRange,
bool isContinuationExpected,
int pageSize,
CancellationToken cancellationToken)
{
requestOptions.MaxItemCount = pageSize;
ResponseMessage message = await this.clientContext.ProcessResourceOperationStreamAsync(
resourceUri : resourceUri,
resourceType : resourceType,
operationType : operationType,
requestOptions : requestOptions,
partitionKey : requestOptions.PartitionKey,
cosmosContainerCore : this.cosmosContainerCore,
streamPayload : this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType),
requestEnricher : (cosmosRequestMessage) =>
{
this.PopulatePartitionKeyRangeInfo(cosmosRequestMessage, partitionKeyRange);
cosmosRequestMessage.Headers.Add(
HttpConstants.HttpHeaders.IsContinuationExpected,
isContinuationExpected.ToString());
QueryRequestOptions.FillContinuationToken(
cosmosRequestMessage,
continuationToken);
cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.ContentType, MediaTypes.QueryJson);
cosmosRequestMessage.Headers.Add(HttpConstants.HttpHeaders.IsQuery, bool.TrueString);
},
diagnosticsContext : null,
cancellationToken : cancellationToken);
QueryResponseCore queryResponseCore = this.GetCosmosElementResponse(
requestOptions,
resourceType,
message,
partitionKeyRange,
queryPageDiagnostics);
if (queryResponseCore.IsSuccess &&
this.clientContext.ClientOptions.EncryptionKeyWrapProvider != null)
{
return(await this.GetDecryptedElementResponseAsync(queryResponseCore, message, cancellationToken));
}
return(queryResponseCore);
}
private static QueryResponseCore CreateFromCosmosException(CosmosException cosmosException)
{
QueryResponseCore queryResponseCore = QueryResponseCore.CreateFailure(
statusCode: cosmosException.StatusCode,
subStatusCodes: (Microsoft.Azure.Documents.SubStatusCodes)cosmosException.SubStatusCode,
errorMessage: cosmosException.ToString(),
requestCharge: 0,
activityId: cosmosException.ActivityId,
diagnostics: QueryResponseCore.EmptyDiagnostics);
return(queryResponseCore);
}
private static QueryResponseCore CreateFromDocumentClientException(Microsoft.Azure.Documents.DocumentClientException documentClientException)
{
QueryResponseCore queryResponseCore = QueryResponseCore.CreateFailure(
statusCode: documentClientException.StatusCode.GetValueOrDefault(System.Net.HttpStatusCode.InternalServerError),
subStatusCodes: null,
errorMessage: documentClientException.ToString(),
requestCharge: 0,
activityId: documentClientException.ActivityId,
diagnostics: QueryResponseCore.EmptyDiagnostics);
return(queryResponseCore);
}
public static QueryResponseCore CreateFromException(Exception exception)
{
QueryResponseCore queryResponseCore;
if (exception is CosmosException cosmosException)
{
queryResponseCore = CreateFromCosmosException(cosmosException);
}
else if (exception is Microsoft.Azure.Documents.DocumentClientException documentClientException)
{
queryResponseCore = CreateFromDocumentClientException(documentClientException);
}
else if (exception is QueryException queryException)
{
CosmosException convertedException = queryException.Accept(QueryExceptionConverter.Singleton);
queryResponseCore = CreateFromCosmosException(convertedException);
}
else if (exception is ExceptionWithStackTraceException exceptionWithStackTrace)
{
queryResponseCore = QueryResponseFactory.CreateFromExceptionWithStackTrace(exceptionWithStackTrace);
}
else
{
if (exception.InnerException != null)
{
// retry with the inner exception
queryResponseCore = QueryResponseFactory.CreateFromException(exception.InnerException);
}
else
{
CosmosException unkownCosmosException = CosmosExceptionFactory.CreateInternalServerErrorException(
headers: new Headers()
{
SubStatusCode = SubStatusCodes.PartitionKeyRangeGone,
ActivityId = QueryResponseCore.EmptyGuidString
},
message: exception.Message,
stackTrace: exception.StackTrace,
trace: NoOpTrace.Singleton,
innerException: exception);
// Unknown exception type should become a 500
queryResponseCore = QueryResponseCore.CreateFailure(
statusCode: System.Net.HttpStatusCode.InternalServerError,
subStatusCodes: null,
cosmosException: unkownCosmosException,
requestCharge: 0,
activityId: QueryResponseCore.EmptyGuidString);
}
}
return(queryResponseCore);
}
public static QueryResponseCore CreateSplitResponse(string collectionRid)
{
QueryResponseCore splitResponse = QueryResponseCore.CreateFailure(
statusCode: HttpStatusCode.Gone,
subStatusCodes: SubStatusCodes.PartitionKeyRangeGone,
errorMessage: "Partition split error",
requestCharge: 10.4,
activityId: Guid.NewGuid().ToString(),
queryMetricsText: null,
queryMetrics: null);
return(splitResponse);
}
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));
}
}
/// <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 async Task <DocumentFeedResponse <CosmosElement> > ExecuteNextFeedResponseAsync(CancellationToken token)
{
QueryResponseCore feedResponse = await this.ExecuteNextAsync(token);
return(new DocumentFeedResponse <CosmosElement>(
result: feedResponse.CosmosElements,
count: feedResponse.CosmosElements.Count,
responseHeaders: new DictionaryNameValueCollection(),
useETagAsContinuation: false,
queryMetrics: null,
requestStats: null,
disallowContinuationTokenMessage: feedResponse.DisallowContinuationTokenMessage,
responseLengthBytes: feedResponse.ResponseLengthBytes));
}
