DATAMONGO-1447 - Add support for isolations on Update.

We now allow usage of the $isolated update operator via Update.isolated().
In case isolated is set the query involved in MongoOperations.updateMulti will be enhanced by '$isolated' : 1 in case the isolation level has not already been set explicitly via eg. new BasicQuery("{'$isolated' : 0}").

Original pull request: spring-projects#371.

Author: christophstrobl

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

@@ -1192,6 +1192,10 @@ public UpdateResult doInCollection(MongoCollection<Document> collection)
 				Document updateObj = update == null ? new Document()
 						: updateMapper.getMappedObject(update.getUpdateObject(), entity);
 
+				if (multi && update.isIsolated() && !queryObj.containsKey("$isolated")) {
+					queryObj.put("$isolated", 1);
+				}
+
 				if (LOGGER.isDebugEnabled()) {
 					LOGGER.debug("Calling update using query: {} and update: {} in collection: {}",
 							serializeToJsonSafely(queryObj), serializeToJsonSafely(updateObj), collectionName);

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

@@ -53,6 +53,7 @@ public enum Position {
 		LAST, FIRST
 	}
 
+	private boolean isolated = false;
 	private Set<String> keysToUpdate = new HashSet<String>();
 	private Map<String, Object> modifierOps = new LinkedHashMap<String, Object>();
 	private Map<String, PushOperatorBuilder> pushCommandBuilders = new LinkedHashMap<String, PushOperatorBuilder>(1);
@@ -73,7 +74,7 @@ public static Update update(String key, Object value) {
 	 * {@literal $set}. This means fields not given in the {@link Document} will be nulled when executing the update. To
 	 * create an only-updating {@link Update} instance of a {@link Document}, call {@link #set(String, Object)} for each
 	 * value in it.
-	 * 
+	 *
 	 * @param object the source {@link Document} to create the update from.
 	 * @param exclude the fields to exclude.
 	 * @return
@@ -364,6 +365,28 @@ public BitwiseOperatorBuilder bitwise(String key) {
 		return new BitwiseOperatorBuilder(this, key);
 	}
 
+	/**
+	 * Prevents a write operation that affects <strong>multiple</strong> documents from yielding to other reads or writes
+	 * once the first document is written. <br />
+	 * Use with {@link org.springframework.data.mongodb.core.MongoOperations#updateMulti(Query, Update, Class)}.
+	 *
+	 * @return never {@literal null}.
+	 * @since 2.0
+	 */
+	public Update isolated() {
+
+		isolated = true;
+		return this;
+	}
+
+	/**
+	 * @return {@literal true} if update isolated is set.
+	 * @since 2.0
+	 */
+	public Boolean isIsolated() {
+		return isolated;
+	}
+
 	public Document getUpdateObject() {
 		return new Document(modifierOps);
 	}

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

@@ -18,6 +18,7 @@
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
+import static org.springframework.data.mongodb.test.util.IsBsonObject.*;
 
 import java.math.BigInteger;
 import java.util.Collections;
@@ -26,6 +27,7 @@
 import java.util.regex.Pattern;
 
 import org.bson.Document;
+import org.bson.conversions.Bson;
 import org.bson.types.ObjectId;
 import org.hamcrest.collection.IsIterableContainingInOrder;
 import org.hamcrest.core.Is;
@@ -337,7 +339,7 @@ public void processDocument(Document document) throws MongoException, DataAccess
 	public void aggregateShouldHonorReadPreferenceWhenSet() {
 
 		when(db.runCommand(Mockito.any(org.bson.Document.class), Mockito.any(ReadPreference.class), eq(Document.class)))
-				.thenReturn(mock(Document.class));
+		.thenReturn(mock(Document.class));
 		template.setReadPreference(ReadPreference.secondary());
 
 		template.aggregate(Aggregation.newAggregation(Aggregation.unwind("foo")), "collection-1", Wrapper.class);
@@ -361,7 +363,7 @@ public void aggregateShouldIgnoreReadPreferenceWhenNotSet() {
 	public void geoNearShouldHonorReadPreferenceWhenSet() {
 
 		when(db.runCommand(Mockito.any(org.bson.Document.class), Mockito.any(ReadPreference.class), eq(Document.class)))
-				.thenReturn(mock(Document.class));
+		.thenReturn(mock(Document.class));
 		template.setReadPreference(ReadPreference.secondary());
 
 		NearQuery query = NearQuery.near(new Point(1, 1));
@@ -374,7 +376,8 @@ public void geoNearShouldHonorReadPreferenceWhenSet() {
 	@Test // DATAMONGO-1166
 	public void geoNearShouldIgnoreReadPreferenceWhenNotSet() {
 
-		when(db.runCommand(Mockito.any(Document.class), eq(Document.class))).thenReturn(mock(Document.class));
+		when(db.runCommand(Mockito.any(Document.class), eq(Document.class))).thenReturn(
+				mock(Document.class));
 
 		NearQuery query = NearQuery.near(new Point(1, 1));
 		template.geoNear(query, Wrapper.class);
@@ -514,6 +517,102 @@ public void onBeforeConvert(BeforeConvertEvent<VersionedEntity> event) {
 		spy.save(entity);
 	}
 
+	@Test // DATAMONGO-1447
+	public void shouldNotAppend$isolatedToNonMulitUpdate() {
+
+		template.updateFirst(new Query(), new Update().isolated().set("jon", "snow"), Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateOne(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().notContaining("$isolated"));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
+	@Test // DATAMONGO-1447
+	public void shouldAppend$isolatedToUpdateMultiEmptyQuery() {
+
+		template.updateMulti(new Query(), new Update().isolated().set("jon", "snow"), Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateMany(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().withSize(1).containing("$isolated", 1));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
+	@Test // DATAMONGO-1447
+	public void shouldAppend$isolatedToUpdateMultiQueryIfNotPresentAndUpdateSetsValue() {
+
+		Update update = new Update().isolated().set("jon", "snow");
+		Query query = new BasicQuery("{'eddard':'stark'}");
+
+		template.updateMulti(query, update, Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateMany(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().containing("$isolated", 1).containing("eddard", "stark"));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
+	@Test // DATAMONGO-1447
+	public void shouldNotAppend$isolatedToUpdateMultiQueryIfNotPresentAndUpdateDoesNotSetValue() {
+
+		Update update = new Update().set("jon", "snow");
+		Query query = new BasicQuery("{'eddard':'stark'}");
+
+		template.updateMulti(query, update, Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateMany(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().notContaining("$isolated").containing("eddard", "stark"));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
+	@Test // DATAMONGO-1447
+	public void shouldNotOverwrite$isolatedToUpdateMultiQueryIfPresentAndUpdateDoesNotSetValue() {
+
+		Update update = new Update().set("jon", "snow");
+		Query query = new BasicQuery("{'eddard':'stark', '$isolated' : 1}");
+
+		template.updateMulti(query, update, Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateMany(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().containing("$isolated", 1).containing("eddard", "stark"));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
+	@Test // DATAMONGO-1447
+	public void shouldNotOverwrite$isolatedToUpdateMultiQueryIfPresentAndUpdateSetsValue() {
+
+		Update update = new Update().isolated().set("jon", "snow");
+		Query query = new BasicQuery("{'eddard':'stark', '$isolated' : 0}");
+
+		template.updateMulti(query, update, Wrapper.class);
+
+		ArgumentCaptor<Bson> queryCaptor = ArgumentCaptor.forClass(Bson.class);
+		ArgumentCaptor<Bson> updateCaptor = ArgumentCaptor.forClass(Bson.class);
+
+		verify(collection).updateMany(queryCaptor.capture(), updateCaptor.capture(), any());
+
+		assertThat(queryCaptor.getValue(), isBsonObject().containing("$isolated", 0).containing("eddard", "stark"));
+		assertThat(updateCaptor.getValue(), isBsonObject().containing("$set.jon", "snow").notContaining("$isolated"));
+	}
+
 	class AutogenerateableId {
 
 		@Id BigInteger id;

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/IsBsonObject.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015 the original author or authors.
+ * Copyright 2015-2017 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,7 +35,8 @@
  */
 public class IsBsonObject<T extends Bson> extends TypeSafeMatcher<T> {
 
-	private List<ExpectedBsonContent> expectations = new ArrayList<ExpectedBsonContent>();;
+	private List<ExpectedBsonContent> expectations = new ArrayList<>();
+	private Integer expectedSize;
 
 	public static <T extends Bson> IsBsonObject<T> isBsonObject() {
 		return new IsBsonObject<T>();
@@ -49,22 +50,33 @@ protected void describeMismatchSafely(T item, Description mismatchDescription) {
 	@Override
 	public void describeTo(Description description) {
 
+		if (expectedSize != null) {
+			description.appendText(String.format("Expected to contain %s  fields. ", expectedSize));
+		}
+
 		for (ExpectedBsonContent expectation : expectations) {
 
 			if (expectation.not) {
-				description.appendText(String.format("Path %s should not be present.", expectation.path));
+				description.appendText(String.format("Path %s should not be present. ", expectation.path));
 			} else if (expectation.value == null) {
-				description.appendText(String.format("Expected to find path %s.", expectation.path));
+				description.appendText(String.format("Expected to find path %s. ", expectation.path));
 			} else {
-				description.appendText(String.format("Expected to find %s for path %s.", expectation.value, expectation.path));
+				description.appendText(String.format("Expected to find %s for path %s. ", expectation.value, expectation.path));
 			}
 		}
-
 	}
 
 	@Override
 	protected boolean matchesSafely(T item) {
 
+		if (expectedSize != null && item instanceof Document) {
+
+			Document document = (Document) item;
+			if (expectedSize != document.keySet().size()) {
+				return false;
+			}
+		}
+
 		if (expectations.isEmpty()) {
 			return true;
 		}
@@ -147,6 +159,12 @@ public IsBsonObject<T> notContaining(String key) {
 		return this;
 	}
 
+	public IsBsonObject<T> withSize(int size) {
+
+		this.expectedSize = Integer.valueOf(size);
+		return this;
+	}
+
 	static class ExpectedBsonContent {
 		String path;
 		Class<?> type;