From b4bc95ce5f3bf50fb55781742a9ff3147de9dc93 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@pivotal.io>
Date: Tue, 16 Apr 2019 09:48:52 +0200
Subject: [PATCH] DATAMONGO-2259 - Add MongoDB 4.2 expanded format to 'out'
 aggregation operation.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

OutOperation now supports the expanded format for the $out aggregation operation if additional parameters, next to the target collection, are given.

    Aggregation.out("out-col").insertDocuments().in("database-2").uniqueKey("field-1“);

    {
        $out : {
            to : "out-col",
            mode : "insertDocuments",
            db : "database-2",
            uniqueKey : "field-1"
        }
    }

We’ll stick to the 2.6 format if only a collection name has been set.

   Aggregation.out("out-col“);

    { $out : "out-col" }

Original pull request: #740.
---
 .../core/aggregation/OutOperation.java        | 209 +++++++++++++++++-
 .../data/mongodb/util/BsonUtils.java          |  18 ++
 .../aggregation/OutOperationUnitTest.java     |  54 +++++
 3 files changed, 277 insertions(+), 4 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/OutOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/OutOperation.java
index fda480f83..a8e99ba33 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/OutOperation.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/OutOperation.java
@@ -15,8 +15,13 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import java.util.Collection;
+
 import org.bson.Document;
+import org.springframework.data.mongodb.util.BsonUtils;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 /**
  * Encapsulates the {@code $out}-operation.
@@ -26,18 +31,158 @@ import org.springframework.util.Assert;
  *
  * @author Nikolay Bogdanov
  * @author Christoph Strobl
- * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/out/">MongoDB Aggregation Framework: $out</a>
+ * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/out/">MongoDB Aggregation Framework:
+ *      $out</a>
  */
 public class OutOperation implements AggregationOperation {
 
+	private final @Nullable String databaseName;
 	private final String collectionName;
+	private final @Nullable Document uniqueKey;
+	private final @Nullable OutMode mode;
 
 	/**
 	 * @param outCollectionName Collection name to export the results. Must not be {@literal null}.
 	 */
 	public OutOperation(String outCollectionName) {
-		Assert.notNull(outCollectionName, "Collection name must not be null!");
-		this.collectionName = outCollectionName;
+		this(null, outCollectionName, null, null);
+
+	}
+
+	/**
+	 * @param databaseName Optional database name the target collection is located in. Can be {@literal null}.
+	 * @param collectionName Collection name to export the results. Must not be {@literal null}. Can be {@literal null}.
+	 * @param uniqueKey Optional unique key spec identify a document in the to collection for replacement or merge.
+	 * @param mode The mode for merging the aggregation pipeline output with the target collection. Can be
+	 *          {@literal null}. {@literal null}.
+	 * @since 2.2
+	 */
+	private OutOperation(@Nullable String databaseName, String collectionName, @Nullable Document uniqueKey,
+			@Nullable OutMode mode) {
+
+		Assert.notNull(collectionName, "Collection name must not be null!");
+
+		this.databaseName = databaseName;
+		this.collectionName = collectionName;
+		this.uniqueKey = uniqueKey;
+		this.mode = mode;
+	}
+
+	/**
+	 * Optionally specify the database of the target collection.
+	 *
+	 * @param database can be {@literal null}. Defaulted to aggregation target database.
+	 * @return new instance of {@link OutOperation}.
+	 * @since 2.2
+	 */
+	public OutOperation in(@Nullable String database) {
+		return new OutOperation(database, collectionName, uniqueKey, mode);
+	}
+
+	/**
+	 * Optionally specify the field that uniquely identify a document in the target collection. <br />
+	 * For convenience the given {@literal key} can either be a single field name or the Json representation of a key
+	 * {@link Document}.
+	 *
+	 * <pre>
+	 *     <code>
+	 *
+	 * // {
+	 * //    "field-1" : 1
+	 * // }
+	 * .uniqueKey("field-1")
+	 *
+	 * // {
+	 * //    "field-1" : 1,
+	 * //    "field-2" : 1
+	 * // }
+	 * .uniqueKey("{ 'field-1' : 1, 'field-2' : 1}")
+	 *
+	 * </code>
+	 * </pre>
+	 *
+	 * @param key can be {@literal null}. Server uses {@literal _id} when {@literal null}.
+	 * @return new instance of {@link OutOperation}.
+	 * @since 2.2
+	 */
+	public OutOperation uniqueKey(@Nullable String key) {
+
+		Document uniqueKey = key == null ? null : BsonUtils.toDocumentOrElse(key, it -> new Document(it, 1));
+		return new OutOperation(databaseName, collectionName, uniqueKey, mode);
+	}
+
+	/**
+	 * Optionally specify the fields that uniquely identify a document in the target collection. <br />
+	 * 
+	 * <pre>
+	 *     <code>
+	 *
+	 * // {
+	 * //    "field-1" : 1
+	 * //    "field-2" : 1
+	 * // }
+	 * .uniqueKeyOf("field-1", "field-2")
+	 *         </code>
+	 * </pre>
+	 *
+	 * @param fields must not be {@literal null}.
+	 * @return new instance of {@link OutOperation}.
+	 * @since 2.2
+	 */
+	public OutOperation uniqueKeyOf(Collection<String> fields) {
+
+		Assert.notNull(fields, "Fields must not be null!");
+
+		Document uniqueKey = new Document();
+		fields.forEach(it -> uniqueKey.append(it, 1));
+
+		return new OutOperation(databaseName, collectionName, uniqueKey, mode);
+	}
+
+	/**
+	 * Specify how to merge the aggregation output with the target collection.
+	 *
+	 * @param mode must not be {@literal null}.
+	 * @return new instance of {@link OutOperation}.
+	 * @since 2.2
+	 */
+	public OutOperation mode(OutMode mode) {
+
+		Assert.notNull(mode, "Mode must not be null!");
+		return new OutOperation(databaseName, collectionName, uniqueKey, mode);
+	}
+
+	/**
+	 * Replace the target collection.
+	 *
+	 * @return new instance of {@link OutOperation}.
+	 * @see OutMode#REPLACE_COLLECTION
+	 * @since 2.2
+	 */
+	public OutOperation replaceCollection() {
+		return mode(OutMode.REPLACE_COLLECTION);
+	}
+
+	/**
+	 * Replace/Upsert documents in the target collection.
+	 *
+	 * @return new instance of {@link OutOperation}.
+	 * @see OutMode#REPLACE
+	 * @since 2.2
+	 */
+	public OutOperation replaceDocuments() {
+		return mode(OutMode.REPLACE);
+	}
+
+	/**
+	 * Insert documents to the target collection.
+	 *
+	 * @return new instance of {@link OutOperation}.
+	 * @see OutMode#INSERT
+	 * @since 2.2
+	 */
+	public OutOperation insertDocuments() {
+		return mode(OutMode.INSERT);
 	}
 
 	/*
@@ -46,6 +191,62 @@ public class OutOperation implements AggregationOperation {
 	 */
 	@Override
 	public Document toDocument(AggregationOperationContext context) {
-		return new Document("$out", collectionName);
+
+		if (!requiresMongoDb42expandedFormat()) {
+			return new Document("$out", collectionName);
+		}
+
+		Assert.state(mode != null, "Mode must not be null!");
+
+		Document $out = new Document("to", collectionName) //
+				.append("mode", mode.getMongoMode());
+
+		if (StringUtils.hasText(databaseName)) {
+			$out.append("db", databaseName);
+		}
+
+		if (uniqueKey != null) {
+			$out.append("uniqueKey", uniqueKey);
+		}
+
+		return new Document("$out", $out);
+	}
+
+	private boolean requiresMongoDb42expandedFormat() {
+		return StringUtils.hasText(databaseName) || mode != null || uniqueKey != null;
+	}
+
+	/**
+	 * The mode for merging the aggregation pipeline output.
+	 *
+	 * @author Christoph Strobl
+	 * @since 2.2
+	 */
+	public enum OutMode {
+
+		/**
+		 * Write documents to the target collection. Errors if a document same uniqueKey already exists.
+		 */
+		INSERT("insertDocuments"),
+
+		/**
+		 * Update on any document in the target collection with the same uniqueKey.
+		 */
+		REPLACE("replaceDocuments"),
+
+		/**
+		 * Replaces the to collection with the output from the aggregation pipeline. Cannot be in a different database.
+		 */
+		REPLACE_COLLECTION("replaceCollection");
+
+		private String mode;
+
+		OutMode(String mode) {
+			this.mode = mode;
+		}
+
+		public String getMongoMode() {
+			return mode;
+		}
 	}
 }
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java
index 56ea876f9..34c85cbc7 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java
@@ -18,12 +18,14 @@ package org.springframework.data.mongodb.util;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.Map;
+import java.util.function.Function;
 
 import org.bson.BsonValue;
 import org.bson.Document;
 import org.bson.conversions.Bson;
 import org.springframework.lang.Nullable;
 import org.springframework.util.ObjectUtils;
+import org.springframework.util.StringUtils;
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
@@ -129,4 +131,20 @@ public class BsonUtils {
 		Arrays.asList(documents).forEach(target::putAll);
 		return target;
 	}
+
+	/**
+	 * @param source
+	 * @param orElse
+	 * @return
+	 * @since 2.2
+	 */
+	public static Document toDocumentOrElse(String source, Function<String, Document> orElse) {
+
+		if (StringUtils.trimLeadingWhitespace(source).startsWith("{")) {
+			return Document.parse(source);
+		}
+
+		return orElse.apply(source);
+	}
+
 }
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/OutOperationUnitTest.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/OutOperationUnitTest.java
index 997fbf864..2f02631fd 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/OutOperationUnitTest.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/OutOperationUnitTest.java
@@ -15,12 +15,19 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import static org.assertj.core.api.Assertions.*;
+import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
+
+import java.util.Arrays;
+
+import org.bson.Document;
 import org.junit.Test;
 
 /**
  * Unit tests for {@link OutOperation}.
  *
  * @author Nikolay Bogdanov
+ * @author Christoph Strobl
  */
 public class OutOperationUnitTest {
 
@@ -28,4 +35,51 @@ public class OutOperationUnitTest {
 	public void shouldCheckNPEInCreation() {
 		new OutOperation(null);
 	}
+
+	@Test // DATAMONGO-2259
+	public void shouldUsePreMongoDB42FormatWhenOnlyCollectionIsPresent() {
+		assertThat(out("out-col").toDocument(Aggregation.DEFAULT_CONTEXT)).isEqualTo(new Document("$out", "out-col"));
+	}
+
+	@Test // DATAMONGO-2259
+	public void shouldUseMongoDB42ExtendedFormatWhenAdditionalParametersPresent() {
+
+		assertThat(out("out-col").insertDocuments().toDocument(Aggregation.DEFAULT_CONTEXT))
+				.isEqualTo(new Document("$out", new Document("to", "out-col").append("mode", "insertDocuments")));
+	}
+
+	@Test // DATAMONGO-2259
+	public void shouldRenderExtendedFormatWithJsonStringKey() {
+
+		assertThat(out("out-col").insertDocuments().in("database-2").uniqueKey("{ 'field-1' : 1, 'field-2' : 1}")
+				.toDocument(Aggregation.DEFAULT_CONTEXT))
+						.isEqualTo(new Document("$out", new Document("to", "out-col").append("mode", "insertDocuments")
+								.append("db", "database-2").append("uniqueKey", new Document("field-1", 1).append("field-2", 1))));
+	}
+
+	@Test // DATAMONGO-2259
+	public void shouldRenderExtendedFormatWithSingleFieldKey() {
+
+		assertThat(
+				out("out-col").insertDocuments().in("database-2").uniqueKey("field-1").toDocument(Aggregation.DEFAULT_CONTEXT))
+						.isEqualTo(new Document("$out", new Document("to", "out-col").append("mode", "insertDocuments")
+								.append("db", "database-2").append("uniqueKey", new Document("field-1", 1))));
+	}
+
+	@Test // DATAMONGO-2259
+	public void shouldRenderExtendedFormatWithMultiFieldKey() {
+
+		assertThat(out("out-col").insertDocuments().in("database-2").uniqueKeyOf(Arrays.asList("field-1", "field-2"))
+				.toDocument(Aggregation.DEFAULT_CONTEXT))
+						.isEqualTo(new Document("$out", new Document("to", "out-col").append("mode", "insertDocuments")
+								.append("db", "database-2").append("uniqueKey", new Document("field-1", 1).append("field-2", 1))));
+	}
+
+	@Test // DATAMONGO-2259
+	public void shouldErrorOnExtendedFormatWithoutMode() {
+
+		assertThatThrownBy(() -> out("out-col").in("database-2").toDocument(Aggregation.DEFAULT_CONTEXT))
+				.isInstanceOf(IllegalStateException.class);
+	}
+
 }