DATADOC-102 modified Update to allow multiple field updates for most operations

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

@@ -15,10 +15,11 @@
  */
 package org.springframework.data.document.mongodb.query;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 
+import org.springframework.dao.InvalidDataAccessApiUsageException;
+
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
 
@@ -28,7 +29,7 @@ public class Update {
 		LAST, FIRST
 	}
 
-  private HashMap<String, Object> criteria = new LinkedHashMap<String, Object>();
+	private HashMap<String, Object> modifierOps = new LinkedHashMap<String, Object>();
 
 	/**
 	 * Static factory method to create an Update using the provided key
@@ -48,7 +49,7 @@ public class Update {
 	 * @return
 	 */
 	public Update set(String key, Object value) {
-    criteria.put("$set", Collections.singletonMap(key, convertValueIfNecessary(value)));
+		addMultiFieldOperation("$set", key, convertValueIfNecessary(value));
 		return this;
 	}
 
@@ -59,7 +60,7 @@ public class Update {
 	 * @return
 	 */
 	public Update unset(String key) {
-    criteria.put("$unset", Collections.singletonMap(key, 1));
+		addMultiFieldOperation("$unset", key, 1);
 		return this;
 	}
 
@@ -71,7 +72,7 @@ public class Update {
 	 * @return
 	 */
 	public Update inc(String key, Number inc) {
-    criteria.put("$inc", Collections.singletonMap(key, inc));
+		addMultiFieldOperation("$inc", key, inc);
 		return this;
 	}
 
@@ -83,7 +84,7 @@ public class Update {
 	 * @return
 	 */
 	public Update push(String key, Object value) {
-    criteria.put("$push", Collections.singletonMap(key, convertValueIfNecessary(value)));
+		addMultiFieldOperation("$push", key, convertValueIfNecessary(value));
 		return this;
 	}
 
@@ -101,7 +102,7 @@ public class Update {
 		}
 		DBObject keyValue = new BasicDBObject();
 		keyValue.put(key, convertedValues);
-    criteria.put("$pushAll", keyValue);
+		modifierOps.put("$pushAll", keyValue);
 		return this;
 	}
 
@@ -113,7 +114,7 @@ public class Update {
 	 * @return
 	 */
 	public Update addToSet(String key, Object value) {
-    criteria.put("$addToSet", Collections.singletonMap(key, convertValueIfNecessary(value)));
+		addMultiFieldOperation("$addToSet", key, convertValueIfNecessary(value));
 		return this;
 	}
 
@@ -125,7 +126,8 @@ public class Update {
 	 * @return
 	 */
 	public Update pop(String key, Position pos) {
-    criteria.put("$pop", Collections.singletonMap(key, (pos == Position.FIRST ? -1 : 1)));
+		addMultiFieldOperation("$pop", key, 
+						(pos == Position.FIRST ? -1 : 1));
 		return this;
 	}
 
@@ -137,7 +139,7 @@ public class Update {
 	 * @return
 	 */
 	public Update pull(String key, Object value) {
-    criteria.put("$pull", Collections.singletonMap(key, convertValueIfNecessary(value)));
+		addMultiFieldOperation("$pull", key, convertValueIfNecessary(value));
 		return this;
 	}
 
@@ -155,7 +157,7 @@ public class Update {
 		}
 		DBObject keyValue = new BasicDBObject();
 		keyValue.put(key, convertedValues);
-    criteria.put("$pullAll", keyValue);
+		modifierOps.put("$pullAll", keyValue);
 		return this;
 	}
 
@@ -167,18 +169,38 @@ public class Update {
 	 * @return
 	 */
 	public Update rename(String oldName, String newName) {
-    criteria.put("$rename", Collections.singletonMap(oldName, newName));
+		addMultiFieldOperation("$rename", oldName, newName);
 		return this;
 	}
 
 	public DBObject getUpdateObject() {
 		DBObject dbo = new BasicDBObject();
-    for (String k : criteria.keySet()) {
-      dbo.put(k, criteria.get(k));
+		for (String k : modifierOps.keySet()) {
+			dbo.put(k, modifierOps.get(k));
 		}
 		return dbo;
 	}
 
+	@SuppressWarnings("unchecked")
+	protected void addMultiFieldOperation(String operator, String key,
+			Object value) {
+		Object existingValue = this.modifierOps.get(operator);
+		LinkedHashMap<String, Object> keyValueMap;
+		if (existingValue == null) {
+			keyValueMap = new LinkedHashMap<String, Object>();
+			this.modifierOps.put(operator, keyValueMap);
+		} else {
+			if (existingValue instanceof LinkedHashMap) {
+				keyValueMap = (LinkedHashMap<String, Object>) existingValue;
+			}
+			else {
+				throw new InvalidDataAccessApiUsageException("Modifier Operations should be a LinkedHashMap but was " + 
+						existingValue.getClass());
+			}
+		}
+		keyValueMap.put(key, value);
+	}
+
 	protected Object convertValueIfNecessary(Object value) {
 		if (value instanceof Enum) {
 			return ((Enum<?>) value).name();

spring-data-mongodb/src/test/java/org/springframework/data/document/mongodb/MongoTemplateTests.java

@@ -15,7 +15,12 @@
  */
 package org.springframework.data.document.mongodb;
 
-import static org.hamcrest.Matchers.*;
+import static org.hamcrest.Matchers.endsWith;
+import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.isOneOf;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.springframework.data.document.mongodb.query.Criteria.where;
 
@@ -23,9 +28,6 @@ import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 
-import com.mongodb.DBCollection;
-import com.mongodb.DBObject;
-import com.mongodb.Mongo;
 import org.bson.types.ObjectId;
 import org.junit.Before;
 import org.junit.Rule;
@@ -37,11 +39,20 @@ import org.springframework.dao.DataIntegrityViolationException;
 import org.springframework.data.document.mongodb.convert.MappingMongoConverter;
 import org.springframework.data.document.mongodb.convert.MongoConverter;
 import org.springframework.data.document.mongodb.mapping.MongoMappingContext;
-import org.springframework.data.document.mongodb.query.*;
+import org.springframework.data.document.mongodb.query.Criteria;
+import org.springframework.data.document.mongodb.query.Index;
 import org.springframework.data.document.mongodb.query.Index.Duplicates;
+import org.springframework.data.document.mongodb.query.Order;
+import org.springframework.data.document.mongodb.query.Query;
+import org.springframework.data.document.mongodb.query.Update;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 
+import com.mongodb.DBCollection;
+import com.mongodb.DBObject;
+import com.mongodb.Mongo;
+import com.mongodb.WriteResult;
+
 /**
  * Integration test for {@link MongoTemplate}.
  *
@@ -339,6 +350,38 @@ public class MongoTemplateTests {
     }
   }
 
+  @Test
+  public void testUsingUpdateWithMultipleSet() throws Exception {
+
+    template.remove(new Query(), PersonWithIdPropertyOfTypeObjectId.class);
+
+    PersonWithIdPropertyOfTypeObjectId p1 = new PersonWithIdPropertyOfTypeObjectId();
+    p1.setFirstName("Sven");
+    p1.setAge(11);
+    template.insert("springdata", p1);
+    PersonWithIdPropertyOfTypeObjectId p2 = new PersonWithIdPropertyOfTypeObjectId();
+    p2.setFirstName("Mary");
+    p2.setAge(21);
+    template.insert("springdata", p2);
+
+    Update u = new Update().set("firstName", "Bob").set("age", 10);
+    
+    WriteResult wr = template.updateMulti("springdata", new Query(), u);
+        
+    assertThat(wr.getN(), is(2));
+    
+    Query q1 = new Query(Criteria.where("age").in(11, 21));
+    List<PersonWithIdPropertyOfTypeObjectId> r1 = template.find("springdata", q1, PersonWithIdPropertyOfTypeObjectId.class);
+    assertThat(r1.size(), is(0));
+    Query q2 = new Query(Criteria.where("age").is(10));
+    List<PersonWithIdPropertyOfTypeObjectId> r2 = template.find("springdata", q2, PersonWithIdPropertyOfTypeObjectId.class);
+    assertThat(r2.size(), is(2));
+    for (PersonWithIdPropertyOfTypeObjectId p : r2) {
+        assertThat(p.getAge(), is(10));
+        assertThat(p.getFirstName(), is("Bob"));
+    }
+  }
+
   @Test
   public void testRemovingDocument() throws Exception {