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

Show more of this group Show more artifacts with this name
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
There is a newer version: 4.61.1
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation.query;

import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.CosmosException;
import com.azure.cosmos.implementation.BadRequestException;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.Resource;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.routing.UInt128;
import com.azure.cosmos.models.FeedResponse;
import com.azure.cosmos.models.ModelBridgeInternal;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;

public class DistinctDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent {
    private final IDocumentQueryExecutionComponent component;
    private final DistinctMap distinctMap;
    private final AtomicReference lastHash;

    private DistinctDocumentQueryExecutionContext(
        IDocumentQueryExecutionComponent component,
        DistinctQueryType distinctQueryType,
        UInt128 previousHash) {
        if (distinctQueryType == DistinctQueryType.NONE) {
            throw new IllegalArgumentException("Invalid distinct query type");
        }

        if (component == null) {
            throw new IllegalArgumentException("documentQueryExecutionComponent cannot be null");
        }

        this.component = component;
        this.distinctMap = DistinctMap.create(distinctQueryType, previousHash);
        this.lastHash = new AtomicReference<>();
    }

    public static  Flux> createAsync(
        BiFunction, Flux>> createSourceComponentFunction,
        DistinctQueryType distinctQueryType,
        String continuationToken,
        PipelinedDocumentQueryParams documentQueryParams) {

        Utils.ValueHolder outDistinctcontinuationtoken = new Utils.ValueHolder<>();
        DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken(null /*lasthash*/,
                                                                                            null /*sourceToken*/);

        if (continuationToken != null) {
            if (!DistinctContinuationToken.tryParse(continuationToken, outDistinctcontinuationtoken)) {
                return Flux.error(new BadRequestException("Invalid DistinctContinuationToken" + continuationToken));
            } else {
                distinctContinuationToken = outDistinctcontinuationtoken.v;
                if (distinctQueryType != DistinctQueryType.ORDERED && distinctContinuationToken.getLastHash() != null) {
                    CosmosException dce = new BadRequestException("DistinctContinuationToken is malformed." +
                                                                              " 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 Flux.error(dce);
                }
            }
        }

        final UInt128 continuationTokenLastHash = distinctContinuationToken.getLastHash();

        return createSourceComponentFunction
            .apply(distinctContinuationToken.getSourceToken(), documentQueryParams)
            .map(component -> new DistinctDocumentQueryExecutionContext(component, distinctQueryType, continuationTokenLastHash));
    }

    IDocumentQueryExecutionComponent getComponent() {
        return this.component;
    }

    @Override
    public Flux> drainAsync(int maxPageSize) {
        return this.component.drainAsync(maxPageSize).map(tFeedResponse -> {
            final List distinctResults = new ArrayList<>();

            tFeedResponse.getResults().forEach(document -> {
                Utils.ValueHolder outHash = new Utils.ValueHolder<>();
                if (this.distinctMap.add(document, outHash)) {
                    distinctResults.add(document);
                    this.lastHash.set(outHash.v);
                }
            });
            Map headers = new HashMap<>(tFeedResponse.getResponseHeaders());
            if (tFeedResponse.getContinuationToken() != null) {

                final String sourceContinuationToken = tFeedResponse.getContinuationToken();
                final DistinctContinuationToken distinctContinuationToken =
                    new DistinctContinuationToken(this.lastHash.get(),
                        sourceContinuationToken);
                headers.put(HttpConstants.HttpHeaders.CONTINUATION,
                            ModelBridgeInternal.toJsonFromJsonSerializable(distinctContinuationToken));
            }

            return BridgeInternal.createFeedResponseWithQueryMetrics(distinctResults,
                headers,
                BridgeInternal.queryMetricsFromFeedResponse(tFeedResponse),
                ModelBridgeInternal.getQueryPlanDiagnosticsContext(tFeedResponse),
                false,
                false,
                tFeedResponse.getCosmosDiagnostics());
        });

    }
}