Support $merge aggregation stage

* Recognize pipelines that end with $merge and treat them as writes
* Recognize the various flavors of the "into" field so that the correct
  find command can be constructed when attempting to iterate an
  aggregation ending with a $merge stage.

JAVA-3319

Author: jyemin

driver-async/src/main/com/mongodb/async/client/AggregateIterable.java

@@ -81,10 +81,10 @@ public interface AggregateIterable<TResult> extends MongoIterable<TResult> {
     AggregateIterable<TResult> useCursor(@Nullable Boolean useCursor);
 
     /**
-     * Aggregates documents according to the specified aggregation pipeline, which must end with a $out stage.
+     * Aggregates documents according to the specified aggregation pipeline, which must end with a $out or $merge stage.
      *
      * @param callback the callback, which is called when the aggregation completes
-     * @throws IllegalStateException if the pipeline does not end with a $out stage
+     * @throws IllegalStateException if the pipeline does not end with a $out or $merge stage
      * @mongodb.driver.manual aggregation/ Aggregation
      */
     void toCollection(SingleResultCallback<Void> callback);
@@ -101,7 +101,7 @@ public interface AggregateIterable<TResult> extends MongoIterable<TResult> {
     /**
      * Sets the bypass document level validation flag.
      *
-     * <p>Note: This only applies when an $out stage is specified</p>.
+     * <p>Note: This only applies when a $out or $merge stage is specified</p>.
      *
      * @param bypassDocumentValidation If true, allows the write to opt-out of document level validation.
      * @return this

driver-async/src/main/com/mongodb/async/client/AggregateIterableImpl.java

@@ -30,7 +30,7 @@
 import com.mongodb.operation.AsyncReadOperation;
 import com.mongodb.operation.AsyncWriteOperation;
 import org.bson.BsonDocument;
-import org.bson.BsonValue;
+import org.bson.BsonString;
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.conversions.Bson;
 
@@ -84,8 +84,8 @@ class AggregateIterableImpl<TDocument, TResult> extends MongoIterableImpl<TResul
     @Override
     public void toCollection(final SingleResultCallback<Void> callback) {
 
-        if (getOutCollection() == null) {
-            throw new IllegalStateException("The last stage of the aggregation pipeline must be $out");
+        if (getOutNamespace() == null) {
+            throw new IllegalStateException("The last stage of the aggregation pipeline must be $out or $merge");
         }
 
         getExecutor().execute(operations.aggregateToCollection(pipeline, maxTimeMS, allowDiskUse, bypassDocumentValidation, collation, hint,
@@ -152,9 +152,9 @@ public AggregateIterable<TResult> hint(@Nullable final Bson hint) {
     @Override
     @SuppressWarnings("deprecation")
     AsyncReadOperation<AsyncBatchCursor<TResult>> asAsyncReadOperation() {
-        BsonValue outCollection = getOutCollection();
+        MongoNamespace outNamespace = getOutNamespace();
 
-        if (outCollection != null) {
+        if (outNamespace != null) {
             AsyncWriteOperation<Void> aggregateToCollectionOperation =
                     operations.aggregateToCollection(pipeline, maxTimeMS, allowDiskUse, bypassDocumentValidation, collation, hint, comment,
                             aggregationLevel);
@@ -165,8 +165,7 @@ AsyncReadOperation<AsyncBatchCursor<TResult>> asAsyncReadOperation() {
                 findOptions.batchSize(batchSize);
             }
             AsyncReadOperation<AsyncBatchCursor<TResult>> findOperation =
-                    operations.find(new MongoNamespace(namespace.getDatabaseName(), outCollection.asString().getValue()),
-                            new BsonDocument(), resultClass, findOptions);
+                    operations.find(outNamespace, new BsonDocument(), resultClass, findOptions);
 
             return new WriteOperationThenCursorReadOperation<TResult>(aggregateToCollectionOperation, findOperation);
         } else {
@@ -177,12 +176,26 @@ AsyncReadOperation<AsyncBatchCursor<TResult>> asAsyncReadOperation() {
     }
 
     @Nullable
-    private BsonValue getOutCollection() {
+    private MongoNamespace getOutNamespace() {
         if (pipeline.size() == 0) {
             return null;
         }
 
         Bson lastStage = notNull("last stage", pipeline.get(pipeline.size() - 1));
-        return lastStage.toBsonDocument(documentClass, codecRegistry).get("$out");
+        BsonDocument lastStageDocument = lastStage.toBsonDocument(documentClass, codecRegistry);
+        if (lastStageDocument.containsKey("$out")) {
+            return new MongoNamespace(namespace.getDatabaseName(), lastStageDocument.getString("$out").getValue());
+        } else if (lastStageDocument.containsKey("$merge")) {
+            BsonDocument mergeDocument = lastStageDocument.getDocument("$merge");
+            if (mergeDocument.isDocument("into")) {
+                BsonDocument intoDocument = mergeDocument.getDocument("into");
+                return new MongoNamespace(intoDocument.getString("db", new BsonString(namespace.getDatabaseName())).getValue(),
+                        intoDocument.getString("coll").getValue());
+            } else if (mergeDocument.isString("into")) {
+                return new MongoNamespace(namespace.getDatabaseName(), mergeDocument.getString("into").getValue());
+            }
+        }
+
+        return null;
     }
 }

driver-async/src/main/com/mongodb/async/client/MapReduceIterable.java

@@ -167,8 +167,6 @@ public interface MapReduceIterable<TResult> extends MongoIterable<TResult> {
     /**
      * Sets the bypass document level validation flag.
      *
-     * <p>Note: This only applies when an $out stage is specified</p>.
-     *
      * @param bypassDocumentValidation If true, allows the write to opt-out of document level validation.
      * @return this
      * @since 3.2

driver-async/src/main/com/mongodb/async/client/MongoCollection.java

@@ -539,7 +539,7 @@ public interface MongoCollection<TDocument> {
     <TResult> FindIterable<TResult> find(ClientSession clientSession, Bson filter, Class<TResult> resultClass);
 
     /**
-     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out stage, the returned
+     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out or $merge stage, the returned
      * iterable will be a query of the collection that the aggregation was written to.  Note that in this case the pipeline will be
      * executed even if the iterable is never iterated.
      *
@@ -550,7 +550,7 @@ public interface MongoCollection<TDocument> {
     AggregateIterable<TDocument> aggregate(List<? extends Bson> pipeline);
 
     /**
-     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out stage, the returned
+     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out or $merge stage, the returned
      * iterable will be a query of the collection that the aggregation was written to.  Note that in this case the pipeline will be
      * executed even if the iterable is never iterated.
      *
@@ -563,7 +563,7 @@ public interface MongoCollection<TDocument> {
     <TResult> AggregateIterable<TResult> aggregate(List<? extends Bson> pipeline, Class<TResult> resultClass);
 
     /**
-     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out stage, the returned
+     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out or $merge stage, the returned
      * iterable will be a query of the collection that the aggregation was written to.  Note that in this case the pipeline will be
      * executed even if the iterable is never iterated.
      *
@@ -577,7 +577,7 @@ public interface MongoCollection<TDocument> {
     AggregateIterable<TDocument> aggregate(ClientSession clientSession, List<? extends Bson> pipeline);
 
     /**
-     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out stage, the returned
+     * Aggregates documents according to the specified aggregation pipeline.  If the pipeline ends with a $out or $merge stage, the returned
      * iterable will be a query of the collection that the aggregation was written to.  Note that in this case the pipeline will be
      * executed even if the iterable is never iterated.
      *

driver-async/src/test/unit/com/mongodb/async/client/AggregateIterableSpecification.groovy

@@ -104,7 +104,7 @@ class AggregateIterableSpecification extends Specification {
         )
     }
 
-    def 'should build the expected AggregateToCollectionOperation'() {
+    def 'should build the expected AggregateToCollectionOperation for $out'() {
         given:
         def cursor = Stub(AsyncBatchCursor) {
             next(_) >> {
@@ -207,6 +207,155 @@ class AggregateIterableSpecification extends Specification {
                 .comment('this is a comment'))
     }
 
+    def 'should build the expected AggregateToCollectionOperation for $merge'() {
+        given:
+        def cursor = Stub(AsyncBatchCursor) {
+            next(_) >> {
+                it[0].onResult(null, null)
+            }
+        }
+        def executor = new TestOperationExecutor([cursor, cursor, cursor, cursor, cursor, cursor, cursor]);
+        def collectionName = 'collectionName'
+        def collectionNamespace = new MongoNamespace(namespace.getDatabaseName(), collectionName)
+        def pipeline = [new Document('$match', 1), new Document('$merge', new Document('into', collectionName))]
+        def pipelineWithIntoDocument = [new Document('$match', 1), new Document('$merge',
+                new Document('into', new Document('db', 'db2').append('coll', collectionName)))]
+
+        when: 'aggregation includes $merge'
+        new AggregateIterableImpl(null, namespace, Document, Document, codecRegistry, readPreference, readConcern, writeConcern, executor,
+                pipeline, AggregationLevel.COLLECTION, true)
+                .batchSize(99)
+                .maxAwaitTime(99, MILLISECONDS)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .useCursor(true)
+                .collation(collation)
+                .hint(new Document('a', 1))
+                .comment('this is a comment')
+                .into([]) { result, t -> }
+
+        def operation = executor.getReadOperation() as WriteOperationThenCursorReadOperation
+
+        then: 'should use the overrides'
+        expect operation.getAggregateToCollectionOperation(), isTheSameAs(new AggregateToCollectionOperation(namespace,
+                [new BsonDocument('$match', new BsonInt32(1)),
+                 new BsonDocument('$merge', new BsonDocument('into', new BsonString(collectionName)))],
+                writeConcern)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .collation(collation)
+                .hint(new BsonDocument('a', new BsonInt32(1)))
+                .comment('this is a comment'))
+
+        when: 'the subsequent read should have the batchSize set'
+        operation = operation.getReadOperation() as FindOperation
+
+        then: 'should use the correct settings'
+        operation.getNamespace() == collectionNamespace
+        operation.getCollation() == collation
+        operation.getBatchSize() == 99
+        operation.getMaxAwaitTime(MILLISECONDS) == 0
+        operation.getMaxTime(MILLISECONDS) == 0
+
+        when: 'aggregation includes $merge into a different database'
+        new AggregateIterableImpl(null, namespace, Document, Document, codecRegistry, readPreference, readConcern, writeConcern, executor,
+                pipelineWithIntoDocument, AggregationLevel.COLLECTION, false)
+                .batchSize(99)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .useCursor(true)
+                .collation(collation)
+                .hint(new Document('a', 1))
+                .comment('this is a comment')
+                .into([]) { result, t -> }
+
+        operation = executor.getReadOperation() as WriteOperationThenCursorReadOperation
+
+        then: 'should use the overrides'
+        expect operation.getAggregateToCollectionOperation(), isTheSameAs(new AggregateToCollectionOperation(namespace,
+                [new BsonDocument('$match', new BsonInt32(1)),
+                 new BsonDocument('$merge', new BsonDocument('into',
+                         new BsonDocument('db', new BsonString('db2')).append('coll', new BsonString(collectionName))))],
+                writeConcern,
+                AggregationLevel.COLLECTION)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .collation(collation)
+                .hint(new BsonDocument('a', new BsonInt32(1)))
+                .comment('this is a comment')
+        )
+
+        when: 'the subsequent read should have the batchSize set'
+        operation = operation.getReadOperation() as FindOperation
+
+        then: 'should use the correct settings'
+        operation.getNamespace() == new MongoNamespace('db2', collectionName)
+        operation.getBatchSize() == 99
+        operation.getCollation() == collation
+        operation.getMaxAwaitTime(MILLISECONDS) == 0
+        operation.getMaxTime(MILLISECONDS) == 0
+
+        when: 'aggregation includes $out and is at the database level'
+        new AggregateIterableImpl(null, namespace, Document, Document, codecRegistry, readPreference, readConcern, writeConcern, executor,
+                pipeline, AggregationLevel.DATABASE, true)
+                .batchSize(99)
+                .maxAwaitTime(99, MILLISECONDS)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .useCursor(true)
+                .collation(collation)
+                .hint(new Document('a', 1))
+                .comment('this is a comment')
+                .into([]) { result, t -> }
+
+        operation = executor.getReadOperation() as WriteOperationThenCursorReadOperation
+
+        then: 'should use the overrides'
+        expect operation.getAggregateToCollectionOperation(), isTheSameAs(new AggregateToCollectionOperation(namespace,
+                [new BsonDocument('$match', new BsonInt32(1)),
+                 new BsonDocument('$merge', new BsonDocument('into', new BsonString(collectionName)))],
+                writeConcern,
+                AggregationLevel.DATABASE)
+                .maxTime(999, MILLISECONDS)
+                .allowDiskUse(true)
+                .collation(collation)
+                .hint(new BsonDocument('a', new BsonInt32(1)))
+                .comment('this is a comment'))
+
+        when: 'the subsequent read should have the batchSize set'
+        operation = operation.getReadOperation() as FindOperation
+
+        then: 'should use the correct settings'
+        operation.getNamespace() == collectionNamespace
+        operation.getCollation() == collation
+        operation.getBatchSize() == 99
+        operation.getMaxAwaitTime(MILLISECONDS) == 0
+        operation.getMaxTime(MILLISECONDS) == 0
+
+        when: 'toCollection should work as expected'
+        def futureResultCallback = new FutureResultCallback()
+        new AggregateIterableImpl(null, namespace, Document, Document, codecRegistry, readPreference, readConcern, writeConcern, executor,
+                pipeline, AggregationLevel.COLLECTION, true)
+                .allowDiskUse(true)
+                .collation(collation)
+                .hint(new Document('a', 1))
+                .comment('this is a comment')
+                .toCollection(futureResultCallback)
+        futureResultCallback.get()
+
+        operation = executor.getWriteOperation() as AggregateToCollectionOperation
+
+        then:
+        expect operation, isTheSameAs(new AggregateToCollectionOperation(namespace,
+                [new BsonDocument('$match', new BsonInt32(1)),
+                 new BsonDocument('$merge', new BsonDocument('into', new BsonString(collectionName)))],
+                writeConcern)
+                .allowDiskUse(true)
+                .collation(collation)
+                .hint(new BsonDocument('a', new BsonInt32(1)))
+                .comment('this is a comment'))
+    }
+
     def 'should handle exceptions correctly'() {
         given:
         def codecRegistry = fromProviders([new ValueCodecProvider(), new BsonValueCodecProvider()])

driver-core/src/main/com/mongodb/operation/AggregateToCollectionOperation.java

@@ -120,8 +120,6 @@ public AggregateToCollectionOperation(final MongoNamespace namespace, final List
         this.aggregationLevel = notNull("aggregationLevel", aggregationLevel);
 
         isTrueArgument("pipeline is not empty", !pipeline.isEmpty());
-        isTrueArgument("last stage of pipeline contains an output collection",
-                pipeline.get(pipeline.size() - 1).get("$out") != null);
     }
 
     /**
@@ -208,7 +206,7 @@ public Boolean getBypassDocumentValidation() {
     /**
      * Sets the bypass document level validation flag.
      *
-     * <p>Note: This only applies when an $out stage is specified</p>.
+     * <p>Note: This only applies when an $out or $merge stage is specified</p>.
      *
      * @param bypassDocumentValidation If true, allows the write to opt-out of document level validation.
      * @return this

driver-core/src/main/com/mongodb/operation/FindAndReplaceOperation.java

@@ -36,11 +36,11 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionThreeDotTwo;
 import static com.mongodb.operation.CommandOperationHelper.CommandCreator;
 import static com.mongodb.operation.DocumentHelper.putIfNotNull;
 import static com.mongodb.operation.DocumentHelper.putIfNotZero;
 import static com.mongodb.operation.DocumentHelper.putIfTrue;
-import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionThreeDotTwo;
 import static com.mongodb.operation.OperationHelper.validateCollation;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
@@ -261,8 +261,6 @@ public Boolean getBypassDocumentValidation() {
     /**
      * Sets the bypass document level validation flag.
      *
-     * <p>Note: This only applies when an $out stage is specified</p>.
-     *
      * @param bypassDocumentValidation If true, allows the write to opt-out of document level validation.
      * @return this
      * @since 3.2

driver-core/src/main/com/mongodb/operation/FindAndUpdateOperation.java

@@ -38,11 +38,11 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.assertions.Assertions.notNull;
+import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionThreeDotTwo;
 import static com.mongodb.operation.CommandOperationHelper.CommandCreator;
 import static com.mongodb.operation.DocumentHelper.putIfNotNull;
 import static com.mongodb.operation.DocumentHelper.putIfNotZero;
 import static com.mongodb.operation.DocumentHelper.putIfTrue;
-import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionThreeDotTwo;
 import static com.mongodb.operation.OperationHelper.validateCollation;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
@@ -264,8 +264,6 @@ public Boolean getBypassDocumentValidation() {
     /**
      * Sets the bypass document level validation flag.
      *
-     * <p>Note: This only applies when an $out stage is specified</p>.
-     *
      * @param bypassDocumentValidation If true, allows the write to opt-out of document level validation.
      * @return this
      * @since 3.2