com.mongodb.reactivestreams.client.internal.MapReducePublisherImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of mongodb-driver-reactivestreams Show documentation
A Reactive Streams implementation of the MongoDB Java driver
There is a newer version: 5.3.0-beta0
Show newest version
/*
 * Copyright 2008-present MongoDB, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.mongodb.reactivestreams.client.internal;

import com.mongodb.MongoNamespace;
import com.mongodb.ReadPreference;
import com.mongodb.client.cursor.TimeoutMode;
import com.mongodb.client.model.Collation;
import com.mongodb.internal.TimeoutSettings;
import com.mongodb.internal.async.AsyncBatchCursor;
import com.mongodb.internal.async.SingleResultCallback;
import com.mongodb.internal.binding.AsyncReadBinding;
import com.mongodb.internal.binding.AsyncWriteBinding;
import com.mongodb.internal.client.model.FindOptions;
import com.mongodb.internal.operation.AsyncOperations;
import com.mongodb.internal.operation.AsyncReadOperation;
import com.mongodb.internal.operation.AsyncWriteOperation;
import com.mongodb.internal.operation.MapReduceAsyncBatchCursor;
import com.mongodb.internal.operation.MapReduceStatistics;
import com.mongodb.lang.Nullable;
import com.mongodb.reactivestreams.client.ClientSession;
import org.bson.BsonDocument;
import org.bson.conversions.Bson;
import org.reactivestreams.Publisher;

import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import static com.mongodb.ReadPreference.primary;
import static com.mongodb.assertions.Assertions.notNull;

@SuppressWarnings("deprecation")
final class MapReducePublisherImpl extends BatchCursorPublisher implements com.mongodb.reactivestreams.client.MapReducePublisher {

    private final String mapFunction;
    private final String reduceFunction;

    private boolean inline = true;
    private String collectionName;
    private String finalizeFunction;
    private Bson scope;
    private Bson filter;
    private Bson sort;
    private int limit;
    private boolean jsMode;
    private boolean verbose = true;
    private long maxTimeMS;
    private com.mongodb.client.model.MapReduceAction action = com.mongodb.client.model.MapReduceAction.REPLACE;
    private String databaseName;
    private Boolean bypassDocumentValidation;
    private Collation collation;

    MapReducePublisherImpl(
            @Nullable final ClientSession clientSession,
            final MongoOperationPublisher mongoOperationPublisher,
            final String mapFunction,
            final String reduceFunction) {
        super(clientSession, mongoOperationPublisher);
        this.mapFunction = notNull("mapFunction", mapFunction);
        this.reduceFunction = notNull("reduceFunction", reduceFunction);
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher collectionName(final String collectionName) {
        this.collectionName = notNull("collectionName", collectionName);
        this.inline = false;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher finalizeFunction(@Nullable final String finalizeFunction) {
        this.finalizeFunction = finalizeFunction;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher scope(@Nullable final Bson scope) {
        this.scope = scope;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher sort(@Nullable final Bson sort) {
        this.sort = sort;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher filter(@Nullable final Bson filter) {
        this.filter = filter;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher limit(final int limit) {
        this.limit = limit;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher jsMode(final boolean jsMode) {
        this.jsMode = jsMode;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher verbose(final boolean verbose) {
        this.verbose = verbose;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher maxTime(final long maxTime, final TimeUnit timeUnit) {
        notNull("timeUnit", timeUnit);
        this.maxTimeMS = TimeUnit.MILLISECONDS.convert(maxTime, timeUnit);
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher action(final com.mongodb.client.model.MapReduceAction action) {
        this.action = action;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher databaseName(@Nullable final String databaseName) {
        this.databaseName = databaseName;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher batchSize(final int batchSize) {
        super.batchSize(batchSize);
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher bypassDocumentValidation(
            @Nullable final Boolean bypassDocumentValidation) {
        this.bypassDocumentValidation = bypassDocumentValidation;
        return this;
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher timeoutMode(final TimeoutMode timeoutMode) {
        super.timeoutMode(timeoutMode);
        return this;
    }

    @Override
    public Publisher toCollection() {
        if (inline) {
            throw new IllegalStateException("The options must specify a non-inline result");
        }
        return getMongoOperationPublisher().createWriteOperationMono(
                (asyncOperations -> asyncOperations.createTimeoutSettings(maxTimeMS)),
                this::createMapReduceToCollectionOperation,
                getClientSession());
    }

    @Override
    public com.mongodb.reactivestreams.client.MapReducePublisher collation(@Nullable final Collation collation) {
        this.collation = collation;
        return this;
    }

    @Override
    ReadPreference getReadPreference() {
        if (inline) {
            return super.getReadPreference();
        } else {
            return primary();
        }
    }

    @Override
    Function, TimeoutSettings> getTimeoutSettings() {
        return (asyncOperations -> asyncOperations.createTimeoutSettings(maxTimeMS));
    }

    @Override
    AsyncReadOperation> asAsyncReadOperation(final int initialBatchSize) {
        if (inline) {
            // initialBatchSize is ignored for map reduce operations.
            return createMapReduceInlineOperation();
        } else {
            return new VoidWriteOperationThenCursorReadOperation<>(createMapReduceToCollectionOperation(),
                    createFindOperation(initialBatchSize));
        }
    }

    private WrappedMapReduceReadOperation createMapReduceInlineOperation() {
        return new WrappedMapReduceReadOperation<>(getOperations().mapReduce(mapFunction, reduceFunction, finalizeFunction,
                getDocumentClass(), filter, limit, jsMode, scope, sort, verbose, collation));
    }

    private WrappedMapReduceWriteOperation createMapReduceToCollectionOperation() {
        return new WrappedMapReduceWriteOperation(
                getOperations().mapReduceToCollection(databaseName, collectionName, mapFunction, reduceFunction, finalizeFunction, filter,
                        limit, jsMode, scope, sort, verbose, action, bypassDocumentValidation, collation));
    }

    private AsyncReadOperation> createFindOperation(final int initialBatchSize) {
        String dbName = databaseName != null ? databaseName : getNamespace().getDatabaseName();
        FindOptions findOptions = new FindOptions().collation(collation).batchSize(initialBatchSize);
        return getOperations().find(new MongoNamespace(dbName, collectionName), new BsonDocument(), getDocumentClass(), findOptions);
    }

    // this could be inlined, but giving it a name so that it's unit-testable
    static class WrappedMapReduceReadOperation implements AsyncReadOperation> {
        private final AsyncReadOperation> operation;

        WrappedMapReduceReadOperation(final AsyncReadOperation> operation) {
            this.operation = operation;
        }

        AsyncReadOperation> getOperation() {
            return operation;
        }

        @Override
        public void executeAsync(final AsyncReadBinding binding, final SingleResultCallback> callback) {
            operation.executeAsync(binding, callback::onResult);
        }
    }

    static class WrappedMapReduceWriteOperation implements AsyncWriteOperation {
        private final AsyncWriteOperation operation;

        WrappedMapReduceWriteOperation(final AsyncWriteOperation operation) {
            this.operation = operation;
        }

        AsyncWriteOperation getOperation() {
            return operation;
        }

        @Override
        public void executeAsync(final AsyncWriteBinding binding, final SingleResultCallback callback) {
            operation.executeAsync(binding, (result, t) -> callback.onResult(null, t));
        }
    }
}