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

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

import com.azure.cosmos.ChangeFeedOptions;
import com.azure.cosmos.FeedOptions;
import com.azure.cosmos.FeedResponse;
import com.azure.cosmos.Resource;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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 Paginator {

    private final static Logger logger = LoggerFactory.getLogger(Paginator.class);

    public static  Flux> getPaginatedChangeFeedQueryResultAsObservable(
            ChangeFeedOptions feedOptions, BiFunction createRequestFunc,
            Function>> executeFunc, Class resourceType,
            int maxPageSize) {
        return getPaginatedQueryResultAsObservable(feedOptions.getRequestContinuation(), createRequestFunc, executeFunc, resourceType,
                -1, maxPageSize, true);
    }

    public static  Flux> getPaginatedQueryResultAsObservable(
            FeedOptions feedOptions,
            BiFunction createRequestFunc,
            Function>> executeFunc, Class resourceType,
            int maxPageSize) {
        return getPaginatedQueryResultAsObservable(feedOptions.requestContinuation(), createRequestFunc, executeFunc, resourceType,
                -1, maxPageSize);
    }

    public static  Flux> getPaginatedQueryResultAsObservable(
            String continuationToken,
            BiFunction createRequestFunc,
            Function>> executeFunc, Class resourceType,
            int top, int maxPageSize) {
        return getPaginatedQueryResultAsObservable(continuationToken, createRequestFunc, executeFunc, resourceType,
                top, maxPageSize, false);
    }

    private static  Flux> getPaginatedQueryResultAsObservable(
            String continuationToken,
            BiFunction createRequestFunc,
            Function>> executeFunc, Class resourceType,
            int top, int maxPageSize, boolean isChangeFeed) {

        return Flux.defer(() -> {
            Flux>> generate = Flux.generate(() ->
                    new Fetcher<>(createRequestFunc, executeFunc, continuationToken, isChangeFeed, top, maxPageSize),
                    (tFetcher, sink) -> {
                        if (tFetcher.shouldFetchMore()) {
                            Flux> nextPage = tFetcher.nextPage();
                            sink.next(nextPage);
                        } else {
                            logger.debug("No more results");
                            sink.complete();
                        }
                        return tFetcher;
            });

            return generate.flatMapSequential(feedResponseFlux -> feedResponseFlux, 1);
        });
    }
}