com.azure.cosmos.implementation.query.OrderByDocumentProducer Maven / Gradle / Ivy

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;

import com.azure.cosmos.implementation.ClientSideRequestStatistics;
import com.azure.cosmos.implementation.DistinctClientSideRequestStatisticsCollection;
import com.azure.cosmos.implementation.Document;
import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
import com.azure.cosmos.implementation.PartitionKeyRange;
import com.azure.cosmos.implementation.QueryMetrics;
import com.azure.cosmos.implementation.RequestChargeTracker;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult;
import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
import com.azure.cosmos.models.FeedResponse;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.function.Supplier;

class OrderByDocumentProducer extends DocumentProducer {

    private static final ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccessor feedResponseAccessor =
        ImplementationBridgeHelpers.FeedResponseHelper.getFeedResponseAccessor();
    private final OrderbyRowComparer consumeComparer;
    private final Map targetRangeToOrderByContinuationTokenMap;

    OrderByDocumentProducer(
            OrderbyRowComparer consumeComparer,
            IDocumentQueryClient client,
            String collectionResourceId,
            CosmosQueryRequestOptions cosmosQueryRequestOptions,
            TriFunction createRequestFunc,
            Function>> executeRequestFunc,
            FeedRangeEpkImpl feedRange,
            String collectionLink,
            Supplier createRetryPolicyFunc,
            Class resourceType,
            UUID correlatedActivityId,
            int initialPageSize,
            String initialContinuationToken,
            int top,
            Map targetRangeToOrderByContinuationTokenMap,
            Supplier operationContextTextProvider) {
        super(client, collectionResourceId, cosmosQueryRequestOptions, createRequestFunc, executeRequestFunc,
              collectionLink, createRetryPolicyFunc, resourceType, correlatedActivityId, initialPageSize,
              initialContinuationToken,top, feedRange, operationContextTextProvider);
        this.consumeComparer = consumeComparer;
        this.targetRangeToOrderByContinuationTokenMap = targetRangeToOrderByContinuationTokenMap;
    }

    protected Flux produceOnFeedRangeGone(Flux> replacementProducers) {
        return replacementProducers.collectList().flux().flatMap(documentProducers -> {
            RequestChargeTracker tracker = new RequestChargeTracker();
            Map queryMetricsMap = new ConcurrentHashMap<>();
            Collection clientSideRequestStatisticsList = new DistinctClientSideRequestStatisticsCollection();
            return OrderByUtils.orderedMerge(consumeComparer, tracker, documentProducers, queryMetricsMap,
                    targetRangeToOrderByContinuationTokenMap, clientSideRequestStatisticsList)
                    .map(orderByQueryResult -> resultPageFrom(tracker, orderByQueryResult));
        });
    }

    private DocumentProducerFeedResponse resultPageFrom(RequestChargeTracker tracker, OrderByRowResult row) {
        double requestCharge = tracker.getAndResetCharge();
        Map headers = Utils.immutableMapOf(HttpConstants.HttpHeaders.REQUEST_CHARGE, String.valueOf(requestCharge));
        FeedResponse fr = feedResponseAccessor.createFeedResponse(
            Collections.singletonList(row), headers, null);
        return new DocumentProducerFeedResponse(fr, row.getSourceRange());
    }

    protected DocumentProducer createChildDocumentProducerOnSplit(
            PartitionKeyRange targetRange,
            String initialContinuationToken) {

        return new OrderByDocumentProducer(
            consumeComparer,
            client,
            collectionRid,
            cosmosQueryRequestOptions,
            createRequestFunc,
            executeRequestFuncWithRetries,
            new FeedRangeEpkImpl(targetRange.toRange()),
            collectionLink,
            createRetryPolicyFunc,
            resourceType ,
            correlatedActivityId,
            pageSize,
            initialContinuationToken,
            top,
            this.targetRangeToOrderByContinuationTokenMap,
            this.operationContextTextProvider);
    }

}