DATAMONGO-2218 - Add support for replaceOne operation in BulkOperations

Author: iqjumps

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/BulkOperations.java

@@ -135,6 +135,15 @@ enum BulkMode {
 	 */
 	BulkOperations remove(List<Query> removes);
 
+	/**
+	 * Add a single replace operation to the bulk operation.
+	 *
+	 * @param query Update criteria.
+	 * @param document the document to replace, must not be {@literal null}.
+	 * @return the current {@link BulkOperations} instance with the replace added, will never be {@literal null}.
+	 */
+	BulkOperations replaceOne(Query query, Object document);
+
 	/**
 	 * Execute all bulk operations using the default write concern.
 	 *

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/DefaultBulkOperations.java

@@ -15,18 +15,12 @@
  */
 package org.springframework.data.mongodb.core;
 
+import com.mongodb.WriteConcern;
+import com.mongodb.client.model.*;
 import lombok.NonNull;
 import lombok.Value;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Optional;
-import java.util.stream.Collectors;
-
 import org.bson.Document;
 import org.bson.conversions.Bson;
-import org.springframework.dao.DataAccessException;
 import org.springframework.dao.support.PersistenceExceptionTranslator;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.convert.UpdateMapper;
@@ -38,18 +32,11 @@
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
-import com.mongodb.BulkWriteException;
-import com.mongodb.WriteConcern;
-import com.mongodb.client.MongoCollection;
-import com.mongodb.client.model.BulkWriteOptions;
-import com.mongodb.client.model.DeleteManyModel;
-import com.mongodb.client.model.DeleteOneModel;
-import com.mongodb.client.model.DeleteOptions;
-import com.mongodb.client.model.InsertOneModel;
-import com.mongodb.client.model.UpdateManyModel;
-import com.mongodb.client.model.UpdateOneModel;
-import com.mongodb.client.model.UpdateOptions;
-import com.mongodb.client.model.WriteModel;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * Default implementation for {@link BulkOperations}.
@@ -266,6 +253,32 @@ public BulkOperations remove(List<Query> removes) {
 		return this;
 	}
 
+    /*
+     * (non-Javadoc)
+     * @see org.springframework.data.mongodb.core.BulkOperations#replaceOne(org.springframework.data.mongodb.core.query.Query, java.lang.Object)
+     */
+	@Override
+	public BulkOperations replaceOne(Query query, Object document) {
+
+		Assert.notNull(query, "Query must not be null!");
+		Assert.notNull(document, "Document must not be null!");
+
+		ReplaceOptions replaceOptions = new ReplaceOptions();
+		query.getCollation().map(Collation::toMongoCollation).ifPresent(replaceOptions::collation);
+		Bson mappedQuery = getMappedQuery(query.getQueryObject());
+
+		if (document instanceof Document) {
+			models.add(new ReplaceOneModel<>(mappedQuery, (Document) document, replaceOptions));
+			return this;
+		}
+
+		Document sink = new Document();
+		mongoOperations.getConverter().write(document, sink);
+		models.add(new ReplaceOneModel<>(mappedQuery, sink, replaceOptions));
+
+		return this;
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.BulkOperations#executeBulk()
@@ -274,7 +287,7 @@ public BulkOperations remove(List<Query> removes) {
 	public com.mongodb.bulk.BulkWriteResult execute() {
 
 		try {
-			
+
 			return mongoOperations.execute(collectionName, collection -> {
 				return collection.bulkWrite(models.stream().map(this::mapWriteModel).collect(Collectors.toList()), bulkOptions);
 			});

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsIntegrationTests.java

@@ -200,6 +200,31 @@ public void removeUnordered() {
 		testRemove(BulkMode.UNORDERED);
 	}
 
+	@Test // DATAMONGO-2218
+	public void replaceOneOrdered() {
+		testReplaceOne(BulkMode.ORDERED);
+	}
+
+	@Test // DATAMONGO-2218
+	public void replaceOneUnordered() {
+		testReplaceOne(BulkMode.UNORDERED);
+	}
+
+	@Test // DATAMONGO-2218
+	public void replaceOneDoesReplace() {
+
+		insertSomeDocuments();
+
+		com.mongodb.bulk.BulkWriteResult result = createBulkOps(BulkMode.ORDERED).//
+				replaceOne(where("_id", "1"), rawDoc("1", "value2")).//
+				execute();
+
+		assertThat(result, notNullValue());
+		assertThat(result.getMatchedCount(), is(1));
+		assertThat(result.getModifiedCount(), is(1));
+		assertThat(result.getInsertedCount(), is(0));
+	}
+
 	/**
 	 * If working on the same set of documents, only an ordered bulk operation will yield predictable results.
 	 */
@@ -278,6 +303,19 @@ private void testRemove(BulkMode mode) {
 		assertThat(createBulkOps(mode).remove(removes).execute().getDeletedCount(), is(3));
 	}
 
+	private void testReplaceOne(BulkMode mode) {
+
+		BulkOperations bulkOps = createBulkOps(mode);
+
+		insertSomeDocuments();
+
+		Query query = where("_id", "1");
+		Document document = rawDoc("1", "value2");
+		int modifiedCount = bulkOps.replaceOne(query, document).execute().getModifiedCount();
+
+		assertThat(modifiedCount, is(1));
+	}
+
 	private BulkOperations createBulkOps(BulkMode mode) {
 		return createBulkOps(mode, null);
 	}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/DefaultBulkOperationsUnitTests.java

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Optional;
 
+import com.mongodb.client.model.*;
 import org.bson.Document;
 import org.junit.Before;
 import org.junit.Test;
@@ -51,10 +52,6 @@
 
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
-import com.mongodb.client.model.DeleteManyModel;
-import com.mongodb.client.model.UpdateManyModel;
-import com.mongodb.client.model.UpdateOneModel;
-import com.mongodb.client.model.WriteModel;
 
 /**
  * Unit tests for {@link DefaultBulkOperations}.
@@ -133,6 +130,18 @@ public void removeShouldUseCollationWhenPresent() {
 				.isEqualTo(com.mongodb.client.model.Collation.builder().locale("de").build());
 	}
 
+	@Test // DATAMONGO-2218
+	public void replaceOneShouldUseCollationWhenPresent() {
+
+		ops.replaceOne(new BasicQuery("{}").collation(Collation.of("de")), new SomeDomainType()).execute();
+
+		verify(collection).bulkWrite(captor.capture(), any());
+
+		assertThat(captor.getValue().get(0)).isInstanceOf(ReplaceOneModel.class);
+		assertThat(((ReplaceOneModel<Document>) captor.getValue().get(0)).getReplaceOptions().getCollation())
+				.isEqualTo(com.mongodb.client.model.Collation.builder().locale("de").build());
+	}
+
 	@Test // DATAMONGO-1678
 	public void bulkUpdateShouldMapQueryAndUpdateCorrectly() {
 
@@ -156,6 +165,23 @@ public void bulkRemoveShouldMapQueryCorrectly() {
 		assertThat(updateModel.getFilter()).isEqualTo(new Document("first_name", "danerys"));
 	}
 
+	@Test // DATAMONGO-2218
+	public void bulkReplaceOneShouldMapQueryCorrectly() {
+
+		SomeDomainType replacement = new SomeDomainType();
+		replacement.firstName = "Minsu";
+		replacement.lastName = "Kim";
+
+		ops.replaceOne(query(where("firstName").is("danerys")), replacement).execute();
+
+		verify(collection).bulkWrite(captor.capture(), any());
+
+		ReplaceOneModel<Document> updateModel = (ReplaceOneModel<Document>) captor.getValue().get(0);
+		assertThat(updateModel.getFilter()).isEqualTo(new Document("first_name", "danerys"));
+		assertThat(updateModel.getReplacement().getString("first_name")).isEqualTo("Minsu");
+		assertThat(updateModel.getReplacement().getString("lastName")).isEqualTo("Kim");
+	}
+
 	class SomeDomainType {
 
 		@Id String id;