From 87ef66cb12fc0f1441cb442d82717ccfc8a79a5c Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@gopivotal.com>
Date: Wed, 23 Apr 2014 10:53:05 +0200
Subject: [PATCH] DATAMONGO-847 - Allow usage of Query within an Update clause.

In case we detect Query within a value used for an Update value we map the query itself to build the expression to use. This allows to form query statements for e.g. $pull using the same API as for the query itself.

Update update = new Update().pull("list", query(where("value").in("foo", "bar")));

Original Pull Request: #172.
---
 .../mongodb/core/convert/UpdateMapper.java    | 19 ++++++++--
 .../core/convert/UpdateMapperUnitTests.java   | 37 +++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/UpdateMapper.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/UpdateMapper.java
index 21154de2a..6e6e0fb51 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/UpdateMapper.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/UpdateMapper.java
@@ -25,6 +25,7 @@ import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty.PropertyToFieldNameConverter;
+import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update.Modifier;
 import org.springframework.data.mongodb.core.query.Update.Modifiers;
 import org.springframework.data.util.ClassTypeInformation;
@@ -79,10 +80,19 @@ public class UpdateMapper extends QueryMapper {
 			return createMapEntry(field, convertSimpleOrDBObject(rawValue, field.getPropertyEntity()));
 		}
 
-		if (!isUpdateModifier(rawValue)) {
-			return super.getMappedObjectForField(field, getMappedValue(field, rawValue));
+		if (isQuery(rawValue)) {
+			return createMapEntry(field,
+					super.getMappedObject(((Query) rawValue).getQueryObject(), field.getPropertyEntity()));
 		}
 
+		if (isUpdateModifier(rawValue)) {
+			return getMappedUpdateModifier(field, rawValue);
+		}
+
+		return super.getMappedObjectForField(field, getMappedValue(field, rawValue));
+	}
+
+	private Entry<String, Object> getMappedUpdateModifier(Field field, Object rawValue) {
 		Object value = null;
 
 		if (rawValue instanceof Modifier) {
@@ -99,7 +109,6 @@ public class UpdateMapper extends QueryMapper {
 
 			value = modificationOperations;
 		} else {
-
 			throw new IllegalArgumentException(String.format("Unable to map value of type '%s'!", rawValue.getClass()));
 		}
 
@@ -119,6 +128,10 @@ public class UpdateMapper extends QueryMapper {
 		return value instanceof Modifier || value instanceof Modifiers;
 	}
 
+	private boolean isQuery(Object value) {
+		return value instanceof Query;
+	}
+
 	private DBObject getMappedValue(Modifier modifier) {
 
 		Object value = converter.convertToMongoType(modifier.getValue(), ClassTypeInformation.OBJECT);
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
index 6722b82b5..619e86290 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/UpdateMapperUnitTests.java
@@ -41,10 +41,13 @@ import org.springframework.data.mongodb.MongoDbFactory;
 import org.springframework.data.mongodb.core.DBObjectTestUtils;
 import org.springframework.data.mongodb.core.mapping.Field;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.data.mongodb.core.query.Criteria;
+import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 
 import com.mongodb.BasicDBList;
 import com.mongodb.BasicDBObject;
+import com.mongodb.BasicDBObjectBuilder;
 import com.mongodb.DBObject;
 import com.mongodb.DBRef;
 
@@ -432,6 +435,40 @@ public class UpdateMapperUnitTests {
 		assertThat(model, allOf(instanceOf(DBRef.class), IsEqual.<Object> equalTo(expectedDBRef)));
 	}
 
+	/**
+	 * @see DATAMONGO-847
+	 */
+	@Test
+	public void updateMapperConvertsNestedQueryCorrectly() {
+
+		Update update = new Update().pull("list", Query.query(Criteria.where("value").in("foo", "bar")));
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(ParentClass.class));
+
+		DBObject $pull = DBObjectTestUtils.getAsDBObject(mappedUpdate, "$pull");
+		DBObject list = DBObjectTestUtils.getAsDBObject($pull, "aliased");
+		DBObject value = DBObjectTestUtils.getAsDBObject(list, "value");
+		BasicDBList $in = DBObjectTestUtils.getAsDBList(value, "$in");
+
+		assertThat($in, IsIterableContainingInOrder.<Object> contains("foo", "bar"));
+	}
+
+	/**
+	 * @see DATAMONGO-847
+	 */
+	@Test
+	public void updateMapperConvertsPullWithNestedQuerfyOnDBRefCorrectly() {
+
+		Update update = new Update().pull("dbRefAnnotatedList", Query.query(Criteria.where("id").is("1")));
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(DocumentWithDBRefCollection.class));
+
+		DBObject $pull = DBObjectTestUtils.getAsDBObject(mappedUpdate, "$pull");
+		DBObject list = DBObjectTestUtils.getAsDBObject($pull, "dbRefAnnotatedList");
+
+		assertThat(list, equalTo(new BasicDBObjectBuilder().add("_id", "1").get()));
+	}
+
 	@org.springframework.data.mongodb.core.mapping.Document(collection = "DocumentWithReferenceToInterface")
 	static interface DocumentWithReferenceToInterface {