DATAMONGO-1236 - Update now include type hint correctly.

We now use property type information when mapping fields affected by an update in case we do not have proper entity information within the context. This allows more precise type resolution required for determining the need to write type hints for a given property.

Original pull request: #301.

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

@@ -22,6 +22,7 @@ import java.util.Map.Entry;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.mapping.Association;
 import org.springframework.data.mapping.context.MappingContext;
+import org.springframework.data.mongodb.core.mapping.BasicMongoPersistentEntity;
 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;
@@ -66,8 +67,8 @@ public class UpdateMapper extends QueryMapper {
 	 */
 	@Override
 	protected Object delegateConvertToMongoType(Object source, MongoPersistentEntity<?> entity) {
-		return entity == null ? super.delegateConvertToMongoType(source, null)
-				: converter.convertToMongoType(source, getTypeHintForEntity(entity));
+		return entity == null ? super.delegateConvertToMongoType(source, null) : converter.convertToMongoType(source,
+				getTypeHintForEntity(source, entity));
 	}
 
 	/*
@@ -141,18 +142,21 @@ public class UpdateMapper extends QueryMapper {
 		return new BasicDBObject(modifier.getKey(), value);
 	}
 
-	private TypeInformation<?> getTypeHintForEntity(MongoPersistentEntity<?> entity) {
-		return processTypeHintForNestedDocuments(entity.getTypeInformation());
+	private TypeInformation<?> getTypeHintForEntity(Object source, MongoPersistentEntity<?> entity) {
+		return processTypeHintForNestedDocuments(source, entity.getTypeInformation());
 	}
 
-	private TypeInformation<?> processTypeHintForNestedDocuments(TypeInformation<?> info) {
+	private TypeInformation<?> processTypeHintForNestedDocuments(Object source, TypeInformation<?> info) {
 
 		Class<?> type = info.getActualType().getType();
 		if (type.isInterface() || java.lang.reflect.Modifier.isAbstract(type.getModifiers())) {
 			return info;
 		}
-		return NESTED_DOCUMENT;
 
+		if (!type.equals(source.getClass())) {
+			return info;
+		}
+		return NESTED_DOCUMENT;
 	}
 
 	/* 
@@ -196,6 +200,18 @@ public class UpdateMapper extends QueryMapper {
 			this.key = key;
 		}
 
+		@Override
+		@SuppressWarnings({ "rawtypes", "unchecked" })
+		public MongoPersistentEntity<?> getPropertyEntity() {
+
+			MongoPersistentEntity<?> entity = super.getPropertyEntity();
+			if (entity != null || getProperty() == null) {
+				return entity;
+			}
+
+			return new BasicMongoPersistentEntity(getProperty().getTypeInformation());
+		}
+
 		/*
 		 * (non-Javadoc)
 		 * @see org.springframework.data.mongodb.core.convert.QueryMapper.MetadataBackedField#getMappedKey()

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

@@ -25,6 +25,7 @@ import static org.springframework.data.mongodb.test.util.IsBsonObject.*;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 
 import org.hamcrest.Matcher;
 import org.hamcrest.collection.IsIterableContainingInOrder;
@@ -680,8 +681,85 @@ public class UpdateMapperUnitTests {
 				context.getPersistentEntity(DomainTypeWrappingConcreteyTypeHavingListOfInterfaceTypeAttributes.class));
 
 		assertThat(mappedUpdate, isBsonObject().notContaining("$set.concreteTypeWithListAttributeOfInterfaceType._class"));
-		assertThat(mappedUpdate, isBsonObject()
-				.containing("$set.concreteTypeWithListAttributeOfInterfaceType.models.[0]._class", ModelImpl.class.getName()));
+		assertThat(
+				mappedUpdate,
+				isBsonObject().containing("$set.concreteTypeWithListAttributeOfInterfaceType.models.[0]._class",
+						ModelImpl.class.getName()));
+	}
+
+	/**
+	 * @see DATAMONGO-1236
+	 */
+	@Test
+	public void mappingShouldRetainTypeInformationForObjectValues() {
+
+		Update update = new Update().set("value", new NestedDocument("kaladin"));
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(EntityWithObject.class));
+
+		assertThat(mappedUpdate, isBsonObject().containing("$set.value.name", "kaladin"));
+		assertThat(mappedUpdate, isBsonObject().containing("$set.value._class", NestedDocument.class.getName()));
+	}
+
+	/**
+	 * @see DATAMONGO-1236
+	 */
+	@Test
+	public void mappingShouldNotRetainTypeInformationForConcreteValues() {
+
+		Update update = new Update().set("concreteValue", new NestedDocument("shallan"));
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(EntityWithObject.class));
+
+		assertThat(mappedUpdate, isBsonObject().containing("$set.concreteValue.name", "shallan"));
+		assertThat(mappedUpdate, isBsonObject().notContaining("$set.concreteValue._class"));
+	}
+
+	/**
+	 * @see DATAMONGO-1236
+	 */
+	@Test
+	public void mappingShouldRetainTypeInformationForObjectValuesWithAlias() {
+
+		Update update = new Update().set("value", new NestedDocument("adolin"));
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(EntityWithAliasedObject.class));
+
+		assertThat(mappedUpdate, isBsonObject().containing("$set.renamed-value.name", "adolin"));
+		assertThat(mappedUpdate, isBsonObject().containing("$set.renamed-value._class", NestedDocument.class.getName()));
+	}
+
+	/**
+	 * @see DATAMONGO-1236
+	 */
+	@Test
+	public void mappingShouldRetrainTypeInformationWhenValueTypeOfMapDoesNotMatchItsDeclaration() {
+
+		Map<Object, Object> map = Collections.<Object, Object> singletonMap("szeth", new NestedDocument("son-son-vallano"));
+
+		Update update = new Update().set("map", map);
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(EntityWithObjectMap.class));
+
+		assertThat(mappedUpdate, isBsonObject().containing("$set.map.szeth.name", "son-son-vallano"));
+		assertThat(mappedUpdate, isBsonObject().containing("$set.map.szeth._class", NestedDocument.class.getName()));
+	}
+
+	/**
+	 * @see DATAMONGO-1236
+	 */
+	@Test
+	public void mappingShouldNotContainTypeInformationWhenValueTypeOfMapMatchesDeclaration() {
+
+		Map<Object, NestedDocument> map = Collections.<Object, NestedDocument> singletonMap("jasnah", new NestedDocument(
+				"kholin"));
+
+		Update update = new Update().set("concreteMap", map);
+		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(),
+				context.getPersistentEntity(EntityWithObjectMap.class));
+
+		assertThat(mappedUpdate, isBsonObject().containing("$set.concreteMap.jasnah.name", "kholin"));
+		assertThat(mappedUpdate, isBsonObject().notContaining("$set.concreteMap.jasnah._class"));
 	}
 
 	static class DomainTypeWrappingConcreteyTypeHavingListOfInterfaceTypeAttributes {
@@ -889,4 +967,21 @@ public class UpdateMapperUnitTests {
 			this.name = name;
 		}
 	}
+
+	static class EntityWithObject {
+
+		Object value;
+		NestedDocument concreteValue;
+	}
+
+	static class EntityWithAliasedObject {
+
+		@Field("renamed-value") Object value;
+	}
+
+	static class EntityWithObjectMap {
+
+		Map<Object, Object> map;
+		Map<Object, NestedDocument> concreteMap;
+	}
 }