com.azure.cosmos.implementation.query.PipelinedDocumentQueryExecutionContext Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of azure-cosmos Show documentation
Show all versions of azure-cosmos Show documentation
This Package contains Microsoft Azure Cosmos SDK (with Reactive Extension Reactor support) for Azure Cosmos DB SQL API
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;
import com.azure.cosmos.implementation.DiagnosticsClientContext;
import com.azure.cosmos.implementation.DocumentCollection;
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
import com.azure.cosmos.implementation.Document;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
import com.azure.cosmos.models.FeedResponse;
import com.azure.cosmos.models.ModelBridgeInternal;
import com.fasterxml.jackson.databind.JsonNode;
import reactor.core.publisher.Flux;
import java.util.function.BiFunction;
import java.util.function.Function;
/**
* While this class is public, but it is not part of our published public APIs.
* This is meant to be internally used only by our sdk.
*/
public class PipelinedDocumentQueryExecutionContext
extends PipelinedQueryExecutionContextBase {
private static final ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.CosmosQueryRequestOptionsAccessor qryOptAccessor =
ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.getCosmosQueryRequestOptionsAccessor();
private final IDocumentQueryExecutionComponent component;
private PipelinedDocumentQueryExecutionContext(
IDocumentQueryExecutionComponent pipelinedComponent,
int actualPageSize,
QueryInfo queryInfo,
Function factoryMethod) {
super(actualPageSize, queryInfo, factoryMethod);
this.component = pipelinedComponent;
}
private static BiFunction, Flux>> createBaseComponentFunction(
DiagnosticsClientContext diagnosticsClientContext,
IDocumentQueryClient client,
PipelinedDocumentQueryParams initParams,
DocumentCollection collection) {
QueryInfo queryInfo = initParams.getQueryInfo();
CosmosQueryRequestOptions requestOptions = initParams.getCosmosQueryRequestOptions();
BiFunction, Flux>> createBaseComponentFunction;
if (queryInfo.hasOrderBy()) {
createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
CosmosQueryRequestOptions orderByCosmosQueryRequestOptions =
qryOptAccessor.clone(requestOptions);
ModelBridgeInternal.setQueryRequestOptionsContinuationToken(orderByCosmosQueryRequestOptions, continuationToken);
ImplementationBridgeHelpers
.CosmosQueryRequestOptionsHelper
.getCosmosQueryRequestOptionsAccessor()
.setItemFactoryMethod(orderByCosmosQueryRequestOptions, null);
documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
return OrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
};
} else {
createBaseComponentFunction = (continuationToken, documentQueryParams) -> {
CosmosQueryRequestOptions parallelCosmosQueryRequestOptions =
qryOptAccessor.clone(requestOptions);
qryOptAccessor.setItemFactoryMethod(parallelCosmosQueryRequestOptions, null);
ModelBridgeInternal.setQueryRequestOptionsContinuationToken(parallelCosmosQueryRequestOptions, continuationToken);
documentQueryParams.setCosmosQueryRequestOptions(parallelCosmosQueryRequestOptions);
return ParallelDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
};
}
BiFunction, Flux>> createAggregateComponentFunction;
if (queryInfo.hasAggregates() && !queryInfo.hasGroupBy()) {
createAggregateComponentFunction =
(continuationToken, documentQueryParams) ->
AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction,
queryInfo.getAggregates(),
queryInfo.getGroupByAliasToAggregateType(),
queryInfo.getGroupByAliases(),
queryInfo.hasSelectValue(),
continuationToken,
documentQueryParams);
} else {
createAggregateComponentFunction = createBaseComponentFunction;
}
return createAggregateComponentFunction;
}
private static
BiFunction, Flux>>
createDistinctPipelineComponentFunction(
BiFunction, Flux>>
createBaseComponent,
QueryInfo queryInfo) {
if (queryInfo.hasDistinct()) {
return
(continuationToken, documentQueryParams) ->
DistinctDocumentQueryExecutionContext.createAsync(createBaseComponent,
queryInfo.getDistinctQueryType(),
continuationToken,
documentQueryParams);
}
return createBaseComponent;
}
private static BiFunction, Flux>> createPipelineComponentFunction(
DiagnosticsClientContext diagnosticsClientContext,
IDocumentQueryClient client,
PipelinedDocumentQueryParams initParams,
DocumentCollection collection) {
BiFunction, Flux>> createBaseComponentFunction = createBaseComponentFunction(
diagnosticsClientContext, client, initParams, collection);
BiFunction, Flux>> createDistinctComponentFunction;
QueryInfo queryInfo = initParams.getQueryInfo();
createDistinctComponentFunction = createDistinctPipelineComponentFunction(
createBaseComponentFunction,
queryInfo);
BiFunction, Flux>> createGroupByComponentFunction;
if (queryInfo.hasGroupBy()) {
createGroupByComponentFunction =
(continuationToken, documentQueryParams) ->
GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction,
continuationToken,
queryInfo.getGroupByAliasToAggregateType(),
queryInfo.getGroupByAliases(),
queryInfo.hasSelectValue(),
documentQueryParams);
} else{
createGroupByComponentFunction = createDistinctComponentFunction;
}
BiFunction, Flux>>
commonPipeLineComponent = createCommonPipelineComponentFunction(
createGroupByComponentFunction,
queryInfo
);
if (queryInfo.hasDCount()) {
return (continuationToken, documentQueryParams) -> DCountDocumentQueryExecutionContext.createAsync(commonPipeLineComponent,
queryInfo,
continuationToken,
documentQueryParams);
} else {
return commonPipeLineComponent;
}
}
protected static Flux> createAsyncCore(
DiagnosticsClientContext diagnosticsClientContext,
IDocumentQueryClient client,
PipelinedDocumentQueryParams initParams,
int pageSize,
Function factoryMethod,
DocumentCollection collection) {
// Use nested callback pattern to unwrap the continuation token and query params at each level.
BiFunction, Flux>> createPipelineComponentFunction =
createPipelineComponentFunction(diagnosticsClientContext, client, initParams.convertGenericType(Document.class), collection);
QueryInfo queryInfo = initParams.getQueryInfo();
CosmosQueryRequestOptions cosmosQueryRequestOptions = initParams.getCosmosQueryRequestOptions();
return createPipelineComponentFunction
.apply(
ModelBridgeInternal.getRequestContinuationFromQueryRequestOptions(cosmosQueryRequestOptions),
initParams.convertGenericType(Document.class))
.map(c -> new PipelinedDocumentQueryExecutionContext<>(
c, pageSize, queryInfo, factoryMethod));
}
@Override
public Flux> executeAsync() {
return this
.component
.drainAsync(this.actualPageSize)
.map(documentFeedResponse -> ImplementationBridgeHelpers
.FeedResponseHelper
.getFeedResponseAccessor().convertGenericType(
documentFeedResponse,
(document) -> this.factoryMethod.apply(document.getPropertyBag())
)
);
}
}