DATAMONGO-407 - Fixed query mapping for updates using collection references.

When an update clause contained a collection element reference (….$.…) we failed to write the type information of the target value object as the key was not translated into a correct property path correctly. We now strip the reference literals and re-apply them when the mapped key is generated.

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

@@ -23,6 +23,7 @@ import java.util.Set;
 import org.bson.types.ObjectId;
 import org.springframework.core.convert.ConversionException;
 import org.springframework.core.convert.ConversionService;
+import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.PropertyPath;
 import org.springframework.data.mapping.PropertyReferenceException;
@@ -30,6 +31,7 @@ import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mapping.context.PersistentPropertyPath;
 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.util.Assert;
 
@@ -102,7 +104,7 @@ public class QueryMapper {
 				continue;
 			}
 
-			Field field = entity == null ? new Field(key) : new MetadataBackedField(key, entity, mappingContext);
+			Field field = createPropertyField(entity, key, mappingContext);
 
 			Object rawValue = query.get(key);
 			String newKey = field.getMappedKey();
@@ -118,6 +120,17 @@ public class QueryMapper {
 		return result;
 	}
 
+	/**
+	 * @param entity
+	 * @param key
+	 * @param mappingContext
+	 * @return
+	 */
+	protected Field createPropertyField(MongoPersistentEntity<?> entity, String key,
+			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext) {
+		return entity == null ? new Field(key) : new MetadataBackedField(key, entity, mappingContext);
+	}
+
 	/**
 	 * Returns the given {@link DBObject} representing a keyword by mapping the keyword's value.
 	 * 
@@ -400,7 +413,7 @@ public class QueryMapper {
 	 * 
 	 * @author Oliver Gierke
 	 */
-	private static class Field {
+	protected static class Field {
 
 		private static final String ID_KEY = "_id";
 
@@ -477,12 +490,14 @@ public class QueryMapper {
 	 * Extension of {@link DocumentField} to be backed with mapping metadata.
 	 * 
 	 * @author Oliver Gierke
+	 * @author Thomas Darimont
 	 */
-	private static class MetadataBackedField extends Field {
+	protected static class MetadataBackedField extends Field {
 
 		private final MongoPersistentEntity<?> entity;
 		private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 		private final MongoPersistentProperty property;
+		private final PersistentPropertyPath<MongoPersistentProperty> path;
 
 		/**
 		 * Creates a new {@link MetadataBackedField} with the given name, {@link MongoPersistentEntity} and
@@ -502,7 +517,7 @@ public class QueryMapper {
 			this.entity = entity;
 			this.mappingContext = context;
 
-			PersistentPropertyPath<MongoPersistentProperty> path = getPath(name);
+			this.path = getPath(name);
 			this.property = path == null ? null : path.getLeafProperty();
 		}
 
@@ -567,19 +582,33 @@ public class QueryMapper {
 		 */
 		@Override
 		public String getMappedKey() {
-
+			return path == null ? name : path.toDotPath(getPropertyConverter());
-			PersistentPropertyPath<MongoPersistentProperty> path = getPath(name);
-			return path == null ? name : path.toDotPath(MongoPersistentProperty.PropertyToFieldNameConverter.INSTANCE);
 		}
 
-		private PersistentPropertyPath<MongoPersistentProperty> getPath(String name) {
+		/**
+		 * Returns the {@link PersistentPropertyPath} for the given <code>pathExpression</code>.
+		 * 
+		 * @param pathExpression
+		 * @return
+		 */
+		private PersistentPropertyPath<MongoPersistentProperty> getPath(String pathExpression) {
 
 			try {
-				PropertyPath path = PropertyPath.from(name, entity.getTypeInformation());
+				PropertyPath path = PropertyPath.from(pathExpression, entity.getTypeInformation());
 				return mappingContext.getPersistentPropertyPath(path);
 			} catch (PropertyReferenceException e) {
 				return null;
 			}
 		}
+
+		/**
+		 * Return the {@link Converter} to be used to created the mapped key. Default implementation will use
+		 * {@link PropertyToFieldNameConverter}.
+		 * 
+		 * @return
+		 */
+		protected Converter<MongoPersistentProperty, String> getPropertyConverter() {
+			return PropertyToFieldNameConverter.INSTANCE;
+		}
 	}
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2014 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.
@@ -15,12 +15,21 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
+import java.util.Arrays;
+import java.util.Iterator;
+
+import org.springframework.core.convert.converter.Converter;
+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.util.Assert;
 
 /**
  * A subclass of {@link QueryMapper} that retains type information on the mongo types.
  * 
  * @author Thomas Darimont
+ * @author Oliver Gierke
  */
 public class UpdateMapper extends QueryMapper {
 
@@ -49,4 +58,90 @@ public class UpdateMapper extends QueryMapper {
 		return entity == null ? super.delegateConvertToMongoType(source, null) : converter.convertToMongoType(source,
 				entity.getTypeInformation());
 	}
+
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.convert.QueryMapper#createPropertyField(org.springframework.data.mongodb.core.mapping.MongoPersistentEntity, java.lang.String, org.springframework.data.mapping.context.MappingContext)
+	 */
+	@Override
+	protected Field createPropertyField(MongoPersistentEntity<?> entity, String key,
+			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext) {
+
+		return entity == null ? super.createPropertyField(entity, key, mappingContext) : //
+				new MetadataBackedUpdateField(entity, key, mappingContext);
+	}
+
+	/**
+	 * {@link MetadataBackedField} that handles {@literal $} paths inside a field key. We clean up an update key
+	 * containing a {@literal $} before handing it to the super class to make sure property lookups and transformations
+	 * continue to work as expected. We provide a custom property converter to re-applied the cleaned up {@literal $}s
+	 * when constructing the mapped key.
+	 * 
+	 * @author Thomas Darimont
+	 * @author Oliver Gierke
+	 */
+	private static class MetadataBackedUpdateField extends MetadataBackedField {
+
+		private final String key;
+
+		/**
+		 * Creates a new {@link MetadataBackedField} with the given {@link MongoPersistentEntity}, key and
+		 * {@link MappingContext}. We clean up the key before handing it up to the super class to make sure it continues to
+		 * work as expected.
+		 * 
+		 * @param entity must not be {@literal null}.
+		 * @param key must not be {@literal null} or empty.
+		 * @param mappingContext must not be {@literal null}.
+		 */
+		public MetadataBackedUpdateField(MongoPersistentEntity<?> entity, String key,
+				MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext) {
+
+			super(key.replaceAll("\\.\\$", ""), entity, mappingContext);
+			this.key = key;
+		}
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.core.convert.QueryMapper.MetadataBackedField#getPropertyConverter()
+		 */
+		@Override
+		protected Converter<MongoPersistentProperty, String> getPropertyConverter() {
+			return new UpdatePropertyConverter(key);
+		}
+
+		/**
+		 * Special {@link Converter} for {@link MongoPersistentProperty} instances that will concatenate the {@literal $}
+		 * contained in the source update key.
+		 * 
+		 * @author Oliver Gierke
+		 */
+		private static class UpdatePropertyConverter implements Converter<MongoPersistentProperty, String> {
+
+			private final Iterator<String> iterator;
+
+			/**
+			 * Creates a new {@link UpdatePropertyConverter} with the given update key.
+			 * 
+			 * @param updateKey must not be {@literal null} or empty.
+			 */
+			public UpdatePropertyConverter(String updateKey) {
+
+				Assert.hasText(updateKey, "Update key must not be null or empty!");
+
+				this.iterator = Arrays.asList(updateKey.split("\\.")).iterator();
+				this.iterator.next();
+			}
+
+			/* 
+			 * (non-Javadoc)
+			 * @see org.springframework.core.convert.converter.Converter#convert(java.lang.Object)
+			 */
+			@Override
+			public String convert(MongoPersistentProperty property) {
+
+				String mappedName = PropertyToFieldNameConverter.INSTANCE.convert(property);
+				return iterator.hasNext() && iterator.next().equals("$") ? String.format("%s.$", mappedName) : mappedName;
+			}
+		}
+	}
 }

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

@@ -167,6 +167,7 @@ public class MongoTemplateTests {
 		template.dropCollection(ObjectWith3AliasedFieldsAndNestedAddress.class);
 		template.dropCollection(BaseDoc.class);
 		template.dropCollection(ObjectWithEnumValue.class);
+		template.dropCollection(DocumentWithCollection.class);
 	}
 
 	@Test
@@ -2205,6 +2206,40 @@ public class MongoTemplateTests {
 		Assert.assertThat(retrieved.model.value(), equalTo("value2"));
 	}
 
+	/**
+	 * @see DATAMONGO-407
+	 */
+	@Test
+	public void updatesShouldRetainTypeInformationEvenForCollections() {
+
+		DocumentWithCollection doc = new DocumentWithCollection();
+		doc.id = "4711";
+		doc.model = new ArrayList<Model>();
+		doc.model.add(new ModelA("foo"));
+		template.insert(doc);
+
+		Query query = new Query(Criteria.where("id").is(doc.id));
+		query.addCriteria(where("model.value").is("foo"));
+		String newModelValue = "bar";
+		Update update = Update.update("model.$", new ModelA(newModelValue));
+		template.updateFirst(query, update, DocumentWithCollection.class);
+
+		Query findQuery = new Query(Criteria.where("id").is(doc.id));
+		DocumentWithCollection result = template.findOne(findQuery, DocumentWithCollection.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.id, is(doc.id));
+		assertThat(result.model, is(notNullValue()));
+		assertThat(result.model, hasSize(1));
+		assertThat(result.model.get(0).value(), is(newModelValue));
+	}
+
+	static class DocumentWithCollection {
+
+		@Id public String id;
+		public List<Model> model;
+	}
+
 	static interface Model {
 		String value();
 	}
@@ -2221,7 +2256,6 @@ public class MongoTemplateTests {
 		public String value() {
 			return this.value;
 		}
-
 	}
 
 	static class Document {

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

@@ -17,6 +17,7 @@ package org.springframework.data.mongodb.core.convert;
 
 import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
+import static org.springframework.data.mongodb.core.DBObjectTestUtils.*;
 
 import java.util.List;
 
@@ -26,7 +27,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.runners.MockitoJUnitRunner;
 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.Update;
 
@@ -64,8 +65,8 @@ public class UpdateMapperUnitTests {
 		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
 				context.getPersistentEntity(ParentClass.class));
 
-		DBObject push = DBObjectTestUtils.getAsDBObject(mappedObject, "$push");
+		DBObject push = getAsDBObject(mappedObject, "$push");
-		DBObject list = DBObjectTestUtils.getAsDBObject(push, "list");
+		DBObject list = getAsDBObject(push, "aliased");
 
 		assertThat(list.get("_class"), is((Object) ConcreteChildClass.class.getName()));
 	}
@@ -82,7 +83,7 @@ public class UpdateMapperUnitTests {
 		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
 				context.getPersistentEntity(ModelWrapper.class));
 
-		DBObject set = DBObjectTestUtils.getAsDBObject(mappedObject, "$set");
+		DBObject set = getAsDBObject(mappedObject, "$set");
 		DBObject modelDbObject = (DBObject) set.get("model");
 		assertThat(modelDbObject.get("_class"), not(nullValue()));
 	}
@@ -99,7 +100,7 @@ public class UpdateMapperUnitTests {
 		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
 				context.getPersistentEntity(ModelWrapper.class));
 
-		DBObject set = DBObjectTestUtils.getAsDBObject(mappedObject, "$set");
+		DBObject set = getAsDBObject(mappedObject, "$set");
 		assertThat(set.get("_class"), nullValue());
 	}
 
@@ -115,14 +116,67 @@ public class UpdateMapperUnitTests {
 		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
 				context.getPersistentEntity(ModelWrapper.class));
 
-		DBObject set = DBObjectTestUtils.getAsDBObject(mappedObject, "$set");
+		DBObject set = getAsDBObject(mappedObject, "$set");
 		assertThat(set.get("_class"), nullValue());
 	}
 
-	static interface Model {
+	/**
+	 * @see DATAMONGO-407
+	 */
+	@Test
+	public void updateMapperShouldRetainTypeInformationForNestedCollectionElements() {
 
+		Update update = Update.update("list.$", new ConcreteChildClass("42", "bubu"));
+
+		UpdateMapper mapper = new UpdateMapper(converter);
+		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(ParentClass.class));
+
+		DBObject set = getAsDBObject(mappedObject, "$set");
+		DBObject modelDbObject = getAsDBObject(set, "aliased.$");
+		assertThat(modelDbObject.get("_class"), is((Object) ConcreteChildClass.class.getName()));
 	}
 
+	/**
+	 * @see DATAMONGO-407
+	 */
+	@Test
+	public void updateMapperShouldSupportNestedCollectionElementUpdates() {
+
+		Update update = Update.update("list.$.value", "foo").set("list.$.otherValue", "bar");
+
+		UpdateMapper mapper = new UpdateMapper(converter);
+		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(ParentClass.class));
+
+		DBObject set = getAsDBObject(mappedObject, "$set");
+		assertThat(set.get("aliased.$.value"), is((Object) "foo"));
+		assertThat(set.get("aliased.$.otherValue"), is((Object) "bar"));
+	}
+
+	/**
+	 * @see DATAMONGO-407
+	 */
+	@Test
+	public void updateMapperShouldWriteTypeInformationForComplexNestedCollectionElementUpdates() {
+
+		Update update = Update.update("list.$.value", "foo").set("list.$.someObject", new ConcreteChildClass("42", "bubu"));
+
+		UpdateMapper mapper = new UpdateMapper(converter);
+		DBObject mappedObject = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(ParentClass.class));
+
+		DBObject dbo = getAsDBObject(mappedObject, "$set");
+		assertThat(dbo.get("aliased.$.value"), is((Object) "foo"));
+
+		DBObject someObject = getAsDBObject(dbo, "aliased.$.someObject");
+		assertThat(someObject, is(notNullValue()));
+		assertThat(someObject.get("_class"), is((Object) ConcreteChildClass.class.getName()));
+		assertThat(someObject.get("value"), is((Object) "bubu"));
+	}
+
+	static interface Model {}
+
 	static class ModelImpl implements Model {
 		public int value;
 
@@ -138,6 +192,8 @@ public class UpdateMapperUnitTests {
 	static class ParentClass {
 
 		String id;
+
+		@Field("aliased")//
 		List<? extends AbstractChildClass> list;
 
 		public ParentClass(String id, List<? extends AbstractChildClass> list) {
@@ -151,10 +207,13 @@ public class UpdateMapperUnitTests {
 
 		String id;
 		String value;
+		String otherValue;
+		AbstractChildClass someObject;
 
 		public AbstractChildClass(String id, String value) {
 			this.id = id;
 			this.value = value;
+			this.otherValue = "other_" + value;
 		}
 	}