private GroupByDocumentQueryExecutionComponent(
CosmosQueryClient cosmosQueryClient,
IReadOnlyDictionary <string, AggregateOperator?> groupByAliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue,
IDocumentQueryExecutionComponent source)
: base(source)
{
if (cosmosQueryClient == null)
{
throw new ArgumentNullException(nameof(cosmosQueryClient));
}
if (groupByAliasToAggregateType == null)
{
throw new ArgumentNullException(nameof(groupByAliasToAggregateType));
}
if (orderedAliases == null)
{
throw new ArgumentNullException(nameof(orderedAliases));
}
this.cosmosQueryClient = cosmosQueryClient;
this.groupingTable = new Dictionary <UInt192, SingleGroupAggregator>();
this.groupByAliasToAggregateType = groupByAliasToAggregateType;
this.orderedAliases = orderedAliases;
this.hasSelectValue = hasSelectValue;
}
/// <summary>
/// Initializes a new instance of the CosmosCrossPartitionQueryExecutionContext class.
/// </summary>
/// <param name="queryContext">Constructor parameters for the base class.</param>
/// <param name="maxConcurrency">The max concurrency</param>
/// <param name="maxBufferedItemCount">The max buffered item count</param>
/// <param name="maxItemCount">Max item count</param>
/// <param name="moveNextComparer">Comparer used to figure out that document producer tree to serve documents from next.</param>
/// <param name="fetchPrioirtyFunction">The priority function to determine which partition to fetch documents from next.</param>
/// <param name="equalityComparer">Used to determine whether we need to return the continuation token for a partition.</param>
/// <param name="returnResultsInDeterministicOrder">Whether or not to return results in deterministic order.</param>
/// <param name="testSettings">Test settings.</param>
protected CosmosCrossPartitionQueryExecutionContext(
CosmosQueryContext queryContext,
int?maxConcurrency,
int?maxItemCount,
int?maxBufferedItemCount,
IComparer <ItemProducerTree> moveNextComparer,
Func <ItemProducerTree, int> fetchPrioirtyFunction,
IEqualityComparer <CosmosElement> equalityComparer,
bool returnResultsInDeterministicOrder,
TestInjections testSettings)
{
if (moveNextComparer == null)
{
throw new ArgumentNullException(nameof(moveNextComparer));
}
this.queryContext = queryContext ?? throw new ArgumentNullException(nameof(queryContext));
this.queryClient = queryContext.QueryClient ?? throw new ArgumentNullException(nameof(queryContext.QueryClient));
this.itemProducerForest = new PriorityQueue <ItemProducerTree>(moveNextComparer, isSynchronized: true);
this.fetchPrioirtyFunction = fetchPrioirtyFunction ?? throw new ArgumentNullException(nameof(fetchPrioirtyFunction));
this.comparableTaskScheduler = new ComparableTaskScheduler(maxConcurrency.GetValueOrDefault(0));
this.equalityComparer = equalityComparer ?? throw new ArgumentNullException(nameof(equalityComparer));
this.testSettings = testSettings;
this.requestChargeTracker = new RequestChargeTracker();
this.diagnosticsPages = new ConcurrentBag <QueryPageDiagnostics>();
this.actualMaxPageSize = maxItemCount.GetValueOrDefault(ParallelQueryConfig.GetConfig().ClientInternalMaxItemCount);
if (this.actualMaxPageSize < 0)
{
throw new ArgumentOutOfRangeException("actualMaxPageSize should never be less than 0");
}
if (this.actualMaxPageSize > int.MaxValue)
{
throw new ArgumentOutOfRangeException("actualMaxPageSize should never be greater than int.MaxValue");
}
if (maxBufferedItemCount.HasValue)
{
this.actualMaxBufferedItemCount = maxBufferedItemCount.Value;
}
else
{
this.actualMaxBufferedItemCount = ParallelQueryConfig.GetConfig().DefaultMaximumBufferSize;
}
if (this.actualMaxBufferedItemCount < 0)
{
throw new ArgumentOutOfRangeException("actualMaxBufferedItemCount should never be less than 0");
}
if (this.actualMaxBufferedItemCount > int.MaxValue)
{
throw new ArgumentOutOfRangeException("actualMaxBufferedItemCount should never be greater than int.MaxValue");
}
this.CanPrefetch = maxConcurrency.HasValue && maxConcurrency.Value != 0;
this.returnResultsInDeterministicOrder = returnResultsInDeterministicOrder;
}
public static async Task <TakeDocumentQueryExecutionComponent> CreateTopDocumentQueryExecutionComponentAsync(
CosmosQueryClient queryClient,
int topCount,
string continuationToken,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback)
{
if (queryClient == null)
{
throw new ArgumentNullException(nameof(queryClient));
}
TopContinuationToken topContinuationToken;
if (continuationToken != null)
{
topContinuationToken = TopContinuationToken.Parse(queryClient, continuationToken);
}
else
{
topContinuationToken = new TopContinuationToken(topCount, null);
}
if (topContinuationToken.Top > topCount)
{
throw queryClient.CreateBadRequestException($"top count in continuation token: {topContinuationToken.Top} can not be greater than the top count in the query: {topCount}.");
}
return(new TakeDocumentQueryExecutionComponent(
await createSourceCallback(topContinuationToken.SourceToken),
topContinuationToken.Top,
TakeEnum.Top));
}
/// <summary>
/// Creates an DistinctDocumentQueryExecutionComponent
/// </summary>
/// <param name="executionEnvironment">The environment to execute on.</param>
/// <param name="queryClient">The query client</param>
/// <param name="requestContinuation">The continuation token.</param>
/// <param name="createSourceCallback">The callback to create the source to drain from.</param>
/// <param name="distinctQueryType">The type of distinct query.</param>
/// <returns>A task to await on and in return </returns>
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
ExecutionEnvironment executionEnvironment,
CosmosQueryClient queryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
DistinctQueryType distinctQueryType)
{
IDocumentQueryExecutionComponent distinctDocumentQueryExecutionComponent;
switch (executionEnvironment)
{
case ExecutionEnvironment.Client:
distinctDocumentQueryExecutionComponent = await ClientDistinctDocumentQueryExecutionComponent.CreateAsync(
queryClient,
requestContinuation,
createSourceCallback,
distinctQueryType);
break;
case ExecutionEnvironment.Compute:
distinctDocumentQueryExecutionComponent = await ComputeDistinctDocumentQueryExecutionComponent.CreateAsync(
queryClient,
requestContinuation,
createSourceCallback,
distinctQueryType);
break;
default:
throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}.");
}
return(distinctDocumentQueryExecutionComponent);
}
public static async Task <TakeDocumentQueryExecutionComponent> CreateLimitDocumentQueryExecutionComponentAsync(
CosmosQueryClient queryClient,
int limitCount,
string continuationToken,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback)
{
if (queryClient == null)
{
throw new ArgumentNullException(nameof(queryClient));
}
LimitContinuationToken limitContinuationToken;
if (continuationToken != null)
{
limitContinuationToken = LimitContinuationToken.Parse(queryClient, continuationToken);
}
else
{
limitContinuationToken = new LimitContinuationToken(limitCount, null);
}
if (limitContinuationToken.Limit > limitCount)
{
throw queryClient.CreateBadRequestException($"limit count in continuation token: {limitContinuationToken.Limit} can not be greater than the limit count in the query: {limitCount}.");
}
return(new TakeDocumentQueryExecutionComponent(
await createSourceCallback(limitContinuationToken.SourceToken),
limitContinuationToken.Limit,
TakeEnum.Limit));
}
/// <summary>
/// Gets the list of partition key ranges.
/// 1. Check partition key range id
/// 2. Check Partition key
/// 3. Check the effective partition key
/// 4. Get the range from the PartitionedQueryExecutionInfo
/// </summary>
internal static async Task <List <Documents.PartitionKeyRange> > GetTargetPartitionKeyRangesAsync(
CosmosQueryClient queryClient,
string resourceLink,
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
ContainerQueryProperties containerQueryProperties,
IReadOnlyDictionary <string, object> properties)
{
List <Documents.PartitionKeyRange> targetRanges;
if (containerQueryProperties.EffectivePartitionKeyString != null)
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
resourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.EffectivePartitionKeyString);
}
else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString))
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
resourceLink,
containerQueryProperties.ResourceId,
effectivePartitionKeyString);
}
else
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
resourceLink,
containerQueryProperties.ResourceId,
partitionedQueryExecutionInfo.QueryRanges);
}
return(targetRanges);
}
/// <summary>
/// Creates an DistinctDocumentQueryExecutionComponent
/// </summary>
/// <param name="queryClient">The query client</param>
/// <param name="requestContinuation">The continuation token.</param>
/// <param name="createSourceCallback">The callback to create the source to drain from.</param>
/// <param name="distinctQueryType">The type of distinct query.</param>
/// <returns>A task to await on and in return </returns>
public static async Task <DistinctDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient queryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
DistinctQueryType distinctQueryType)
{
DistinctContinuationToken distinctContinuationToken;
if (requestContinuation != null)
{
if (!DistinctContinuationToken.TryParse(requestContinuation, out distinctContinuationToken))
{
throw queryClient.CreateBadRequestException($"Invalid {nameof(DistinctContinuationToken)}: {requestContinuation}");
}
}
else
{
distinctContinuationToken = new DistinctContinuationToken(sourceToken: null, distinctMapToken: null);
}
DistinctMap distinctMap = DistinctMap.Create(distinctQueryType, distinctContinuationToken.DistinctMapToken);
IDocumentQueryExecutionComponent source = await createSourceCallback(distinctContinuationToken.SourceToken);
return(new ComputeDistinctDocumentQueryExecutionComponent(
distinctQueryType,
distinctMap,
source));
}
public static Task <PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsync(
CosmosQueryClient client,
SqlQuerySpec sqlQuerySpec,
Uri resourceLink,
PartitionKey?partitionKey,
CancellationToken cancellationToken = default)
{
if (client == null)
{
throw new ArgumentNullException(nameof(client));
}
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}
if (resourceLink == null)
{
throw new ArgumentNullException(nameof(resourceLink));
}
cancellationToken.ThrowIfCancellationRequested();
return(client.ExecuteQueryPlanRequestAsync(
resourceLink,
ResourceType.Document,
OperationType.QueryPlan,
sqlQuerySpec,
partitionKey,
QueryPlanRetriever.SupportedQueryFeaturesString,
cancellationToken));
}
public static async Task <PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
CosmosQueryClient queryClient,
SqlQuerySpec sqlQuerySpec,
Documents.ResourceType resourceType,
PartitionKeyDefinition partitionKeyDefinition,
bool hasLogicalPartitionKey,
bool useSystemPrefix,
ITrace trace,
CancellationToken cancellationToken = default)
{
if (queryClient == null)
{
throw new ArgumentNullException(nameof(queryClient));
}
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}
if (partitionKeyDefinition == null)
{
throw new ArgumentNullException(nameof(partitionKeyDefinition));
}
cancellationToken.ThrowIfCancellationRequested();
using (ITrace serviceInteropTrace = trace.StartChild("Service Interop Query Plan", TraceComponent.Query, TraceLevel.Info))
{
QueryPlanHandler queryPlanHandler = new QueryPlanHandler(queryClient);
TryCatch <PartitionedQueryExecutionInfo> tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync(
sqlQuerySpec,
resourceType,
partitionKeyDefinition,
QueryPlanRetriever.SupportedQueryFeatures,
hasLogicalPartitionKey,
useSystemPrefix,
cancellationToken);
if (!tryGetQueryPlan.Succeeded)
{
Exception originalException = ExceptionWithStackTraceException.UnWrapMonadExcepion(tryGetQueryPlan.Exception, serviceInteropTrace);
if (originalException is CosmosException)
{
throw originalException;
}
throw CosmosExceptionFactory.CreateBadRequestException(
message: originalException.Message,
headers: new Headers(),
stackTrace: tryGetQueryPlan.Exception.StackTrace,
innerException: originalException,
trace: trace);
}
return(tryGetQueryPlan.Result);
}
}
public NetworkAttachedDocumentContainer(
ContainerCore container,
CosmosQueryClient cosmosQueryClient,
CosmosClientContext cosmosClientContext,
CosmosDiagnosticsContext diagnosticsContext,
QueryRequestOptions queryRequestOptions = null)
{
this.container = container ?? throw new ArgumentNullException(nameof(container));
this.cosmosQueryClient = cosmosQueryClient ?? throw new ArgumentNullException(nameof(cosmosQueryClient));
this.cosmosClientContext = cosmosClientContext ?? throw new ArgumentNullException(nameof(cosmosClientContext));
this.diagnosticsContext = diagnosticsContext;
this.queryRequestOptions = queryRequestOptions;
}
public CosmosLinqQueryProvider(
CosmosContainerCore container,
CosmosJsonSerializer cosmosJsonSerializer,
CosmosQueryClient queryClient,
QueryRequestOptions cosmosQueryRequestOptions,
bool allowSynchronousQueryExecution)
{
this.container = container;
this.cosmosJsonSerializer = cosmosJsonSerializer;
this.queryClient = queryClient;
this.cosmosQueryRequestOptions = cosmosQueryRequestOptions;
this.allowSynchronousQueryExecution = allowSynchronousQueryExecution;
}
/// <summary>
/// Parses the TopContinuationToken from it's string form.
/// </summary>
/// <param name="queryClient">The query client</param>
/// <param name="value">The string form to parse from.</param>
/// <returns>The parsed TopContinuationToken.</returns>
public static TopContinuationToken Parse(CosmosQueryClient queryClient, string value)
{
TopContinuationToken result;
if (!TryParse(value, out result))
{
throw queryClient.CreateBadRequestException($"Invalid TopContinuationToken: {value}");
}
else
{
return(result);
}
}
public static async Task <PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
CosmosQueryClient queryClient,
SqlQuerySpec sqlQuerySpec,
PartitionKeyDefinition partitionKeyDefinition,
bool hasLogicalPartitionKey,
ITrace trace,
CancellationToken cancellationToken = default)
{
if (queryClient == null)
{
throw new ArgumentNullException(nameof(queryClient));
}
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}
if (partitionKeyDefinition == null)
{
throw new ArgumentNullException(nameof(partitionKeyDefinition));
}
cancellationToken.ThrowIfCancellationRequested();
using (ITrace serviceInteropTrace = trace.StartChild("Service Interop Query Plan", TraceComponent.Query, TraceLevel.Info))
{
QueryPlanHandler queryPlanHandler = new QueryPlanHandler(queryClient);
TryCatch <PartitionedQueryExecutionInfo> tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync(
sqlQuerySpec,
partitionKeyDefinition,
QueryPlanRetriever.SupportedQueryFeatures,
hasLogicalPartitionKey,
cancellationToken);
if (!tryGetQueryPlan.Succeeded)
{
if (tryGetQueryPlan.Exception is CosmosException)
{
throw tryGetQueryPlan.Exception;
}
throw CosmosExceptionFactory.CreateBadRequestException(
message: tryGetQueryPlan.Exception.ToString(),
stackTrace: tryGetQueryPlan.Exception.StackTrace);
}
return(tryGetQueryPlan.Result);
}
}
public CosmosLinqQuery(
CosmosContainerCore container,
CosmosJsonSerializer cosmosJsonSerializer,
CosmosQueryClient queryClient,
QueryRequestOptions cosmosQueryRequestOptions,
bool allowSynchronousQueryExecution)
: this(
container,
cosmosJsonSerializer,
queryClient,
cosmosQueryRequestOptions,
null,
allowSynchronousQueryExecution)
{
}
public NetworkAttachedDocumentContainer(
ContainerInternal container,
CosmosQueryClient cosmosQueryClient,
CosmosDiagnosticsContext diagnosticsContext,
QueryRequestOptions queryRequestOptions = null,
string resourceLink = null,
ResourceType resourceType = ResourceType.Document)
{
this.container = container ?? throw new ArgumentNullException(nameof(container));
this.cosmosQueryClient = cosmosQueryClient ?? throw new ArgumentNullException(nameof(cosmosQueryClient));
this.diagnosticsContext = diagnosticsContext;
this.queryRequestOptions = queryRequestOptions;
this.resourceLink = resourceLink ?? this.container.LinkUri;
this.resourceType = resourceType;
}
public static CosmosQueryContext CreateContext(
CosmosQueryClient cosmosQueryClient)
{
return(new CosmosQueryContextCore(
client: cosmosQueryClient,
queryRequestOptions: null,
resourceTypeEnum: ResourceType.Document,
operationType: OperationType.Query,
resourceType: typeof(ToDoItem),
resourceLink: DefaultResourceLink,
correlatedActivityId: Guid.NewGuid(),
isContinuationExpected: true,
allowNonValueAggregateQuery: true,
containerResourceId: DefaultCollectionRid));
}
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient cosmosQueryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
IReadOnlyDictionary <string, AggregateOperator?> groupByAliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue)
{
// We do not support continuation tokens for GROUP BY.
return(new GroupByDocumentQueryExecutionComponent(
cosmosQueryClient,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue,
await createSourceCallback(requestContinuation)));
}
/// <summary>
/// Initializes a new instance of the DistinctDocumentQueryExecutionComponent class.
/// </summary>
/// <param name="queryClient">The query client</param>
/// <param name="distinctQueryType">The type of distinct query.</param>
/// <param name="previousHash">The previous that distinct map saw.</param>
/// <param name="source">The source to drain from.</param>
private DistinctDocumentQueryExecutionComponent(
CosmosQueryClient queryClient,
DistinctQueryType distinctQueryType,
UInt192?previousHash,
IDocumentQueryExecutionComponent source)
: base(source)
{
if (distinctQueryType == DistinctQueryType.None)
{
throw new ArgumentException("It doesn't make sense to create a distinct component of type None.");
}
this.queryClient = queryClient;
this.distinctQueryType = distinctQueryType;
this.distinctMap = DistinctMap.Create(distinctQueryType, previousHash);
}
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient cosmosQueryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
IReadOnlyDictionary <string, AggregateOperator?> groupByAliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue)
{
GroupByContinuationToken groupByContinuationToken;
if (requestContinuation != null)
{
if (!GroupByContinuationToken.TryParse(requestContinuation, out groupByContinuationToken))
{
throw cosmosQueryClient.CreateBadRequestException(
$"Invalid {nameof(GroupByContinuationToken)}: '{requestContinuation}'");
}
}
else
{
groupByContinuationToken = new GroupByContinuationToken(
groupingTableContinuationToken: null,
sourceContinuationToken: null);
}
IDocumentQueryExecutionComponent source;
if (groupByContinuationToken.SourceContinuationToken == ComputeGroupByDocumentQueryExecutionComponent.DoneReadingGroupingsContinuationToken)
{
source = DoneDocumentQueryExecutionComponent.Value;
}
else
{
source = await createSourceCallback(groupByContinuationToken.SourceContinuationToken);
}
GroupingTable groupingTable = GroupingTable.CreateFromContinuationToken(
cosmosQueryClient,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue,
groupByContinuationToken.GroupingTableContinuationToken);
return(new ComputeGroupByDocumentQueryExecutionComponent(
source,
groupingTable));
}
public static GroupingTable CreateFromContinuationToken(CosmosQueryClient cosmosQueryClient,
IReadOnlyDictionary <string, AggregateOperator?> groupByAliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue,
string groupingTableContinuationToken)
{
GroupingTable groupingTable = new GroupingTable(
cosmosQueryClient,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue);
if (groupingTableContinuationToken != null)
{
if (!CosmosElement.TryParse(
groupingTableContinuationToken,
out CosmosObject parsedGroupingTableContinuations))
{
throw cosmosQueryClient.CreateBadRequestException($"Invalid GroupingTableContinuationToken");
}
foreach (KeyValuePair <string, CosmosElement> kvp in parsedGroupingTableContinuations)
{
string key = kvp.Key;
CosmosElement value = kvp.Value;
UInt128 groupByKey = UInt128.Parse(key);
if (!(value is CosmosString singleGroupAggregatorContinuationToken))
{
throw cosmosQueryClient.CreateBadRequestException($"Invalid GroupingTableContinuationToken");
}
SingleGroupAggregator singleGroupAggregator = SingleGroupAggregator.Create(
cosmosQueryClient,
EmptyAggregateOperators,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue,
singleGroupAggregatorContinuationToken.Value);
groupingTable.table[groupByKey] = singleGroupAggregator;
}
}
return(groupingTable);
}
public NetworkAttachedDocumentContainer(
ContainerInternal container,
CosmosQueryClient cosmosQueryClient,
Guid correlatedActivityId,
QueryRequestOptions queryRequestOptions = null,
ChangeFeedRequestOptions changeFeedRequestOptions = null,
string resourceLink = null,
ResourceType resourceType = ResourceType.Document)
{
this.container = container ?? throw new ArgumentNullException(nameof(container));
this.cosmosQueryClient = cosmosQueryClient ?? throw new ArgumentNullException(nameof(cosmosQueryClient));
this.queryRequestOptions = queryRequestOptions;
this.changeFeedRequestOptions = changeFeedRequestOptions;
this.resourceLink = resourceLink ?? this.container.LinkUri;
this.resourceType = resourceType;
this.correlatedActivityId = correlatedActivityId;
}
/// <summary>
/// Creates a AggregateDocumentQueryExecutionComponent.
/// </summary>
/// <param name="executionEnvironment">The environment to execute on.</param>
/// <param name="queryClient">The query client.</param>
/// <param name="aggregates">The aggregates.</param>
/// <param name="aliasToAggregateType">The alias to aggregate type.</param>
/// <param name="orderedAliases">The ordering of the aliases.</param>
/// <param name="hasSelectValue">Whether or not the query has the 'VALUE' keyword.</param>
/// <param name="requestContinuation">The continuation token to resume from.</param>
/// <param name="createSourceCallback">The callback to create the source component that supplies the local aggregates.</param>
/// <returns>The AggregateDocumentQueryExecutionComponent.</returns>
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
ExecutionEnvironment executionEnvironment,
CosmosQueryClient queryClient,
AggregateOperator[] aggregates,
IReadOnlyDictionary <string, AggregateOperator?> aliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback)
{
IDocumentQueryExecutionComponent aggregateDocumentQueryExecutionComponent;
switch (executionEnvironment)
{
case ExecutionEnvironment.Client:
aggregateDocumentQueryExecutionComponent = await ClientAggregateDocumentQueryExecutionComponent.CreateAsync(
queryClient,
aggregates,
aliasToAggregateType,
orderedAliases,
hasSelectValue,
requestContinuation,
createSourceCallback);
break;
case ExecutionEnvironment.Compute:
aggregateDocumentQueryExecutionComponent = await ComputeAggregateDocumentQueryExecutionComponent.CreateAsync(
queryClient,
aggregates,
aliasToAggregateType,
orderedAliases,
hasSelectValue,
requestContinuation,
createSourceCallback);
break;
default:
throw new ArgumentException($"Unknown {nameof(ExecutionEnvironment)}: {executionEnvironment}.");
}
return(aggregateDocumentQueryExecutionComponent);
}
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient queryClient,
AggregateOperator[] aggregates,
IReadOnlyDictionary <string, AggregateOperator?> aliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback)
{
string sourceContinuationToken;
string singleGroupAggregatorContinuationToken;
if (requestContinuation != null)
{
if (!AggregateContinuationToken.TryParse(requestContinuation, out AggregateContinuationToken aggregateContinuationToken))
{
throw queryClient.CreateBadRequestException($"Malfomed {nameof(AggregateContinuationToken)}: '{requestContinuation}'");
}
sourceContinuationToken = aggregateContinuationToken.SourceContinuationToken;
singleGroupAggregatorContinuationToken = aggregateContinuationToken.SingleGroupAggregatorContinuationToken;
}
else
{
sourceContinuationToken = null;
singleGroupAggregatorContinuationToken = null;
}
IDocumentQueryExecutionComponent source = await createSourceCallback(sourceContinuationToken);
SingleGroupAggregator singleGroupAggregator = SingleGroupAggregator.Create(
queryClient,
aggregates,
aliasToAggregateType,
orderedAliases,
hasSelectValue,
singleGroupAggregatorContinuationToken);
return(new ComputeAggregateDocumentQueryExecutionComponent(
source,
singleGroupAggregator,
hasSelectValue));
}
public static async Task <PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
CosmosQueryClient queryClient,
SqlQuerySpec sqlQuerySpec,
PartitionKeyDefinition partitionKeyDefinition,
bool hasLogicalPartitionKey,
CancellationToken cancellationToken = default)
{
if (queryClient == null)
{
throw new ArgumentNullException(nameof(queryClient));
}
if (sqlQuerySpec == null)
{
throw new ArgumentNullException(nameof(sqlQuerySpec));
}
if (partitionKeyDefinition == null)
{
throw new ArgumentNullException(nameof(partitionKeyDefinition));
}
cancellationToken.ThrowIfCancellationRequested();
QueryPlanHandler queryPlanHandler = new QueryPlanHandler(queryClient);
TryCatch <PartitionedQueryExecutionInfo> tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync(
sqlQuerySpec,
partitionKeyDefinition,
QueryPlanRetriever.SupportedQueryFeatures,
hasLogicalPartitionKey,
cancellationToken);
if (!tryGetQueryPlan.Succeeded)
{
throw new CosmosException(
System.Net.HttpStatusCode.BadRequest,
tryGetQueryPlan.Exception.ToString());
}
return(tryGetQueryPlan.Result);
}
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient cosmosQueryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
IReadOnlyDictionary <string, AggregateOperator?> groupByAliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue)
{
IDocumentQueryExecutionComponent source = await createSourceCallback(requestContinuation);
GroupingTable groupingTable = GroupingTable.CreateFromContinuationToken(
cosmosQueryClient,
groupByAliasToAggregateType,
orderedAliases,
hasSelectValue,
groupingTableContinuationToken: null);
return(new ClientGroupByDocumentQueryExecutionComponent(
source,
groupingTable));
}
public CosmosLinqQuery(
CosmosContainerCore container,
CosmosJsonSerializer cosmosJsonSerializer,
CosmosQueryClient queryClient,
QueryRequestOptions cosmosQueryRequestOptions,
Expression expression,
bool allowSynchronousQueryExecution)
{
this.container = container ?? throw new ArgumentNullException(nameof(container));
this.cosmosJsonSerializer = cosmosJsonSerializer;
this.queryClient = queryClient;
this.cosmosQueryRequestOptions = cosmosQueryRequestOptions;
this.expression = expression ?? Expression.Constant(this);
this.allowSynchronousQueryExecution = allowSynchronousQueryExecution;
this.queryProvider = new CosmosLinqQueryProvider(
container,
cosmosJsonSerializer,
queryClient,
cosmosQueryRequestOptions,
this.allowSynchronousQueryExecution);
this.correlatedActivityId = Guid.NewGuid();
}
/// <summary>
/// Creates an DistinctDocumentQueryExecutionComponent
/// </summary>
/// <param name="queryClient">The query client</param>
/// <param name="requestContinuation">The continuation token.</param>
/// <param name="createSourceCallback">The callback to create the source to drain from.</param>
/// <param name="distinctQueryType">The type of distinct query.</param>
/// <returns>A task to await on and in return </returns>
public static async Task <DistinctDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient queryClient,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback,
DistinctQueryType distinctQueryType)
{
DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken(null, null);
if (requestContinuation != null)
{
distinctContinuationToken = DistinctContinuationToken.Parse(queryClient, requestContinuation);
if (distinctQueryType != DistinctQueryType.Ordered && distinctContinuationToken.LastHash != null)
{
throw queryClient.CreateBadRequestException($"DistinctContinuationToken is malformed: {distinctContinuationToken}. DistinctContinuationToken can not have a 'lastHash', when the query type is not ordered (ex SELECT DISTINCT VALUE c.blah FROM c ORDER BY c.blah).");
}
}
return(new DistinctDocumentQueryExecutionComponent(
queryClient,
distinctQueryType,
distinctContinuationToken.LastHash,
await createSourceCallback(distinctContinuationToken.SourceToken)));
}
public static async Task <IDocumentQueryExecutionComponent> CreateAsync(
CosmosQueryClient queryClient,
AggregateOperator[] aggregates,
IReadOnlyDictionary <string, AggregateOperator?> aliasToAggregateType,
IReadOnlyList <string> orderedAliases,
bool hasSelectValue,
string requestContinuation,
Func <string, Task <IDocumentQueryExecutionComponent> > createSourceCallback)
{
IDocumentQueryExecutionComponent source = await createSourceCallback(requestContinuation);
SingleGroupAggregator singleGroupAggregator = SingleGroupAggregator.Create(
queryClient,
aggregates,
aliasToAggregateType,
orderedAliases,
hasSelectValue,
continuationToken: null);
return(new ClientAggregateDocumentQueryExecutionComponent(
source,
singleGroupAggregator,
hasSelectValue));
}
/// <summary>
/// Initializes a new instance of the ItemProducerTree class.
/// </summary>
/// <param name="queryContext">query context.</param>
/// <param name="querySpecForInit">query spec init.</param>
/// <param name="partitionKeyRange">The partition key range.</param>
/// <param name="produceAsyncCompleteCallback">Callback to invoke once a fetch finishes.</param>
/// <param name="itemProducerTreeComparer">Comparer to determine, which tree to produce from.</param>
/// <param name="equalityComparer">Comparer to see if we need to return the continuation token for a partition.</param>
/// <param name="testSettings">Test flags.</param>
/// <param name="deferFirstPage">Whether or not to defer fetching the first page.</param>
/// <param name="collectionRid">The collection to drain from.</param>
/// <param name="initialPageSize">The initial page size.</param>
/// <param name="initialContinuationToken">The initial continuation token.</param>
public ItemProducerTree(
CosmosQueryContext queryContext,
SqlQuerySpec querySpecForInit,
Documents.PartitionKeyRange partitionKeyRange,
ProduceAsyncCompleteDelegate produceAsyncCompleteCallback,
IComparer <ItemProducerTree> itemProducerTreeComparer,
IEqualityComparer <CosmosElement> equalityComparer,
TestInjections testSettings,
bool deferFirstPage,
string collectionRid,
long initialPageSize = 50,
string initialContinuationToken = null)
{
if (queryContext == null)
{
throw new ArgumentNullException($"{nameof(queryContext)}");
}
if (itemProducerTreeComparer == null)
{
throw new ArgumentNullException($"{nameof(itemProducerTreeComparer)}");
}
if (produceAsyncCompleteCallback == null)
{
throw new ArgumentNullException($"{nameof(produceAsyncCompleteCallback)}");
}
if (itemProducerTreeComparer == null)
{
throw new ArgumentNullException($"{nameof(itemProducerTreeComparer)}");
}
if (equalityComparer == null)
{
throw new ArgumentNullException($"{nameof(equalityComparer)}");
}
if (string.IsNullOrEmpty(collectionRid))
{
throw new ArgumentException($"{nameof(collectionRid)} can not be null or empty.");
}
this.Root = new ItemProducer(
queryContext,
querySpecForInit,
partitionKeyRange,
(itemsBuffered, resourceUnitUsage, diagnostics, requestLength, token) => produceAsyncCompleteCallback(this, itemsBuffered, resourceUnitUsage, diagnostics, requestLength, token),
equalityComparer,
testSettings,
initialPageSize,
initialContinuationToken);
this.queryClient = queryContext.QueryClient;
this.children = new PriorityQueue <ItemProducerTree>(itemProducerTreeComparer, true);
this.deferFirstPage = deferFirstPage;
this.collectionRid = collectionRid;
this.createItemProducerTreeCallback = ItemProducerTree.CreateItemProducerTreeCallback(
queryContext,
querySpecForInit,
produceAsyncCompleteCallback,
itemProducerTreeComparer,
equalityComparer,
testSettings,
deferFirstPage,
collectionRid,
initialPageSize);
this.executeWithSplitProofingSemaphore = new SemaphoreSlim(1, 1);
}
private static async Task <TryCatch <CosmosQueryExecutionContext> > TryCreateCoreContextAsync(
CosmosQueryContext cosmosQueryContext,
InputParameters inputParameters,
CancellationToken cancellationToken)
{
// The default
using (cosmosQueryContext.CreateDiagnosticScope("CreateQueryPipeline"))
{
// Try to parse the continuation token.
CosmosElement continuationToken = inputParameters.InitialUserContinuationToken;
PartitionedQueryExecutionInfo queryPlanFromContinuationToken = inputParameters.PartitionedQueryExecutionInfo;
if (continuationToken != null)
{
if (!PipelineContinuationToken.TryCreateFromCosmosElement(
continuationToken,
out PipelineContinuationToken pipelineContinuationToken))
{
return(TryCatch <CosmosQueryExecutionContext> .FromException(
new MalformedContinuationTokenException(
$"Malformed {nameof(PipelineContinuationToken)}: {continuationToken}.")));
}
if (PipelineContinuationToken.IsTokenFromTheFuture(pipelineContinuationToken))
{
return(TryCatch <CosmosQueryExecutionContext> .FromException(
new MalformedContinuationTokenException(
$"{nameof(PipelineContinuationToken)} Continuation token is from a newer version of the SDK. " +
$"Upgrade the SDK to avoid this issue." +
$"{continuationToken}.")));
}
if (!PipelineContinuationToken.TryConvertToLatest(
pipelineContinuationToken,
out PipelineContinuationTokenV1_1 latestVersionPipelineContinuationToken))
{
return(TryCatch <CosmosQueryExecutionContext> .FromException(
new MalformedContinuationTokenException(
$"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")));
}
continuationToken = latestVersionPipelineContinuationToken.SourceContinuationToken;
if (latestVersionPipelineContinuationToken.QueryPlan != null)
{
queryPlanFromContinuationToken = latestVersionPipelineContinuationToken.QueryPlan;
}
}
CosmosQueryClient cosmosQueryClient = cosmosQueryContext.QueryClient;
ContainerQueryProperties containerQueryProperties = await cosmosQueryClient.GetCachedContainerQueryPropertiesAsync(
cosmosQueryContext.ResourceLink,
inputParameters.PartitionKey,
cancellationToken);
cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId;
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo;
if (inputParameters.ForcePassthrough)
{
partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo()
{
QueryInfo = new QueryInfo()
{
Aggregates = null,
DistinctType = DistinctQueryType.None,
GroupByAliases = null,
GroupByAliasToAggregateType = null,
GroupByExpressions = null,
HasSelectValue = false,
Limit = null,
Offset = null,
OrderBy = null,
OrderByExpressions = null,
RewrittenQuery = null,
Top = null,
},
QueryRanges = new List <Documents.Routing.Range <string> >(),
};
}
else if (queryPlanFromContinuationToken != null)
{
partitionedQueryExecutionInfo = queryPlanFromContinuationToken;
}
else
{
// If the query would go to gateway, but we have a partition key,
// then try seeing if we can execute as a passthrough using client side only logic.
// This is to short circuit the need to go to the gateway to get the query plan.
if (cosmosQueryContext.QueryClient.ByPassQueryParsing() &&
inputParameters.PartitionKey.HasValue)
{
bool parsed;
SqlQuery sqlQuery;
using (cosmosQueryContext.CreateDiagnosticScope("QueryParsing"))
{
parsed = SqlQuery.TryParse(inputParameters.SqlQuerySpec.QueryText, out sqlQuery);
}
if (parsed)
{
bool hasDistinct = sqlQuery.SelectClause.HasDistinct;
bool hasGroupBy = sqlQuery.GroupByClause != default;
bool hasAggregates = AggregateProjectionDetector.HasAggregate(sqlQuery.SelectClause.SelectSpec);
bool createPassthroughQuery = !hasAggregates && !hasDistinct && !hasGroupBy;
if (createPassthroughQuery)
{
TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
if (responseStats != null)
{
responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
}
// Only thing that matters is that we target the correct range.
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
List <Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
cosmosQueryContext.ResourceLink,
containerQueryProperties.ResourceId,
inputParameters.PartitionKey.Value.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition));
return(await CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContextAsync(
cosmosQueryContext,
inputParameters,
partitionedQueryExecutionInfo : new PartitionedQueryExecutionInfo(),
targetRanges,
containerQueryProperties.ResourceId,
cancellationToken));
}
}
}
if (cosmosQueryContext.QueryClient.ByPassQueryParsing())
{
// For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case.
// We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this
partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync(
cosmosQueryContext,
inputParameters.SqlQuerySpec,
cosmosQueryContext.ResourceLink,
inputParameters.PartitionKey,
cancellationToken);
}
else
{
using (cosmosQueryContext.CreateDiagnosticScope("ServiceInterop"))
{
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
cosmosQueryContext.QueryClient,
inputParameters.SqlQuerySpec,
partitionKeyDefinition,
inputParameters.PartitionKey != null,
cancellationToken);
}
}
}
return(await TryCreateFromPartitionedQuerExecutionInfoAsync(
partitionedQueryExecutionInfo,
containerQueryProperties,
cosmosQueryContext,
inputParameters,
cancellationToken));
}
}
