DATAMONGO-392 - Updates now retain type information.

Extracted mongo type conversion in QueryMapper into delegateConvertToMongoType(…). Introduced QueryMapper subclass UpdateQueryMapper to retain type information in delegateConvertToMongoType(…). Added test case updatesShouldRetainTypeInformation to MongoTemplateTests.

Original pull request: #49.

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

@@ -57,6 +57,7 @@ import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
 import org.springframework.data.mongodb.core.convert.MongoWriter;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
+import org.springframework.data.mongodb.core.convert.UpdateMapper;
 import org.springframework.data.mongodb.core.geo.Distance;
 import org.springframework.data.mongodb.core.geo.GeoResult;
 import org.springframework.data.mongodb.core.geo.GeoResults;
@@ -135,7 +136,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 	private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 	private final MongoDbFactory mongoDbFactory;
 	private final MongoExceptionTranslator exceptionTranslator = new MongoExceptionTranslator();
-	private final QueryMapper mapper;
+	private final QueryMapper queryMapper;
+	private final UpdateMapper updateMapper;
 
 	private WriteConcern writeConcern;
 	private WriteConcernResolver writeConcernResolver = DefaultWriteConcernResolver.INSTANCE;
@@ -188,7 +190,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 
 		this.mongoDbFactory = mongoDbFactory;
 		this.mongoConverter = mongoConverter == null ? getDefaultMongoConverter(mongoDbFactory) : mongoConverter;
-		this.mapper = new QueryMapper(this.mongoConverter);
+		this.queryMapper = new QueryMapper(this.mongoConverter);
+		this.updateMapper = new UpdateMapper(this.mongoConverter);
 
 		// We always have a mapping context in the converter, whether it's a simple one or not
 		mappingContext = this.mongoConverter.getMappingContext();
@@ -495,7 +498,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 			throw new InvalidDataAccessApiUsageException("Query passed in to exist can't be null");
 		}
 
-		DBObject mappedQuery = mapper.getMappedObject(query.getQueryObject(), getPersistentEntity(entityClass));
+		DBObject mappedQuery = queryMapper.getMappedObject(query.getQueryObject(), getPersistentEntity(entityClass));
 		return execute(collectionName, new FindCallback(mappedQuery)).hasNext();
 	}
 
@@ -604,7 +607,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 	private long count(Query query, Class<?> entityClass, String collectionName) {
 
 		Assert.hasText(collectionName);
-		final DBObject dbObject = query == null ? null : mapper.getMappedObject(query.getQueryObject(),
+		final DBObject dbObject = query == null ? null : queryMapper.getMappedObject(query.getQueryObject(),
 				entityClass == null ? null : mappingContext.getPersistentEntity(entityClass));
 
 		return execute(collectionName, new CollectionCallback<Long>() {
@@ -940,10 +943,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 
 				MongoPersistentEntity<?> entity = entityClass == null ? null : getPersistentEntity(entityClass);
 
-				DBObject queryObj = query == null ? new BasicDBObject()
-						: mapper.getMappedObject(query.getQueryObject(), entity);
-				DBObject updateObj = update == null ? new BasicDBObject() : mapper.getMappedObject(update.getUpdateObject(),
+				DBObject queryObj = query == null ? new BasicDBObject() : queryMapper.getMappedObject(query.getQueryObject(),
 						entity);
+				DBObject updateObj = update == null ? new BasicDBObject() : updateMapper.getMappedObject(
+						update.getUpdateObject(), entity);
 
 				if (LOGGER.isDebugEnabled()) {
 					LOGGER.debug("Calling update using query: " + queryObj + " and update: " + updateObj + " in collection: "
@@ -1061,7 +1064,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 
 				maybeEmitEvent(new BeforeDeleteEvent<T>(queryObject, entityClass));
 
-				DBObject dboq = mapper.getMappedObject(queryObject, entity);
+				DBObject dboq = queryMapper.getMappedObject(queryObject, entity);
 
 				MongoAction mongoAction = new MongoAction(writeConcern, MongoActionOperation.REMOVE, collectionName,
 						entityClass, null, queryObject);
@@ -1156,7 +1159,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 		if (criteria == null) {
 			dbo.put("cond", null);
 		} else {
-			dbo.put("cond", mapper.getMappedObject(criteria.getCriteriaObject(), null));
+			dbo.put("cond", queryMapper.getMappedObject(criteria.getCriteriaObject(), null));
 		}
 		// If initial document was a JavaScript string, potentially loaded by Spring's Resource abstraction, load it and
 		// convert to DBObject
@@ -1326,7 +1329,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 	protected <T> T doFindOne(String collectionName, DBObject query, DBObject fields, Class<T> entityClass) {
 		EntityReader<? super T, DBObject> readerToUse = this.mongoConverter;
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
-		DBObject mappedQuery = mapper.getMappedObject(query, entity);
+		DBObject mappedQuery = queryMapper.getMappedObject(query, entity);
 
 		return executeFindOneInternal(new FindOneCallback(mappedQuery, fields), new ReadDbObjectCallback<T>(readerToUse,
 				entityClass), collectionName);
@@ -1362,7 +1365,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 					serializeToJsonSafely(query), fields, entityClass, collectionName));
 		}
 
-		return executeFindMultiInternal(new FindCallback(mapper.getMappedObject(query, entity), fields), preparer,
+		return executeFindMultiInternal(new FindCallback(queryMapper.getMappedObject(query, entity), fields), preparer,
 				objectCallback, collectionName);
 	}
 
@@ -1383,7 +1386,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 		}
 		EntityReader<? super T, DBObject> readerToUse = this.mongoConverter;
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
-		return executeFindMultiInternal(new FindCallback(mapper.getMappedObject(query, entity), fields), null,
+		return executeFindMultiInternal(new FindCallback(queryMapper.getMappedObject(query, entity), fields), null,
 				new ReadDbObjectCallback<T>(readerToUse, entityClass), collectionName);
 	}
 
@@ -1422,7 +1425,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 					+ entityClass + " in collection: " + collectionName);
 		}
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
-		return executeFindOneInternal(new FindAndRemoveCallback(mapper.getMappedObject(query, entity), fields, sort),
+		return executeFindOneInternal(new FindAndRemoveCallback(queryMapper.getMappedObject(query, entity), fields, sort),
 				new ReadDbObjectCallback<T>(readerToUse, entityClass), collectionName);
 	}
 
@@ -1437,8 +1440,8 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);
 
-		DBObject mappedUpdate = mapper.getMappedObject(update.getUpdateObject(), entity);
-		DBObject mappedQuery = mapper.getMappedObject(query, entity);
+		DBObject mappedUpdate = queryMapper.getMappedObject(update.getUpdateObject(), entity);
+		DBObject mappedQuery = queryMapper.getMappedObject(query, entity);
 
 		if (LOGGER.isDebugEnabled()) {
 			LOGGER.debug("findAndModify using query: " + mappedQuery + " fields: " + fields + " sort: " + sort

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

@@ -204,13 +204,25 @@ public class QueryMapper {
 	private Object convertSimpleOrDBObject(Object source, MongoPersistentEntity<?> entity) {
 
 		if (source instanceof BasicDBList) {
-			return converter.convertToMongoType(source);
+			return delegateConvertToMongoType(source, entity);
 		}
 
 		if (source instanceof DBObject) {
 			return getMappedObject((DBObject) source, entity);
 		}
 
+		return delegateConvertToMongoType(source, entity);
+	}
+
+	/**
+	 * Converts the given source Object to a mongo type with the type information of the original source type omitted.
+	 * Subclasses may overwrite this method to retain the type information of the source type on the resulting mongo type.
+	 * 
+	 * @param source
+	 * @param entity
+	 * @return the converted mongo type or null if source is null
+	 */
+	protected Object delegateConvertToMongoType(Object source, MongoPersistentEntity<?> entity) {
 		return converter.convertToMongoType(source);
 	}
 
@@ -262,7 +274,7 @@ public class QueryMapper {
 			// Ignore
 		}
 
-		return converter.convertToMongoType(id);
+		return delegateConvertToMongoType(id, null);
 	}
 
 	/**

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2013 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.convert;
+
+import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
+
+/**
+ * A subclass of {@link QueryMapper} that retains type information on the mongo types.
+ * 
+ * @author Thomas Darimont
+ */
+public class UpdateMapper extends QueryMapper {
+
+	private final MongoWriter<?> converter;
+
+	/**
+	 * Creates a new {@link UpdateMapper} using the given {@link MongoConverter}.
+	 * 
+	 * @param converter must not be {@literal null}.
+	 */
+	public UpdateMapper(MongoConverter converter) {
+
+		super(converter);
+		this.converter = converter;
+	}
+
+	/**
+	 * Converts the given source object to a mongo type retaining the original type information of the source type on the
+	 * mongo type.
+	 * 
+	 * @see org.springframework.data.mongodb.core.convert.QueryMapper#delegateConvertToMongoType(java.lang.Object,
+	 *      org.springframework.data.mongodb.core.mapping.MongoPersistentEntity)
+	 */
+	@Override
+	protected Object delegateConvertToMongoType(Object source, MongoPersistentEntity<?> entity) {
+		return entity == null ? super.delegateConvertToMongoType(source, null) : converter.convertToMongoType(source,
+				entity.getTypeInformation());
+	}
+}

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

@@ -153,6 +153,7 @@ public class MongoTemplateTests {
 		template.dropCollection(TypeWithDate.class);
 		template.dropCollection("collection");
 		template.dropCollection("personX");
+		template.dropCollection(Document.class);
 	}
 
 	@Test
@@ -1681,6 +1682,58 @@ public class MongoTemplateTests {
 		assertThat(result.id, is(idValue));
 	}
 
+	/**
+	 * @see DATAMONGO-392
+	 */
+	@Test
+	public void updatesShouldRetainTypeInformation() {
+
+		Document doc = new Document();
+		doc.id = "4711";
+		doc.model = new ModelA().withValue("foo");
+		template.insert(doc);
+
+		Query query = new Query(Criteria.where("id").is(doc.id));
+		String newModelValue = "bar";
+		Update update = Update.update("model", new ModelA().withValue(newModelValue));
+		template.updateFirst(query, update, Document.class);
+
+		Document result = template.findOne(query, Document.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.id, is(doc.id));
+		assertThat(result.model, is(notNullValue()));
+		assertThat(result.model.value(), is(newModelValue));
+	}
+
+	static interface Model {
+		String value();
+
+		Model withValue(String value);
+	}
+
+	static class ModelA implements Model {
+
+		private String value;
+
+		@Override
+		public String value() {
+			return this.value;
+		}
+
+		@Override
+		public Model withValue(String value) {
+			this.value = value;
+			return this;
+		}
+	}
+
+	static class Document {
+
+		@Id public String id;
+		public Model model;
+	}
+
 	static class MyId {
 
 		String first;
@@ -1692,7 +1745,7 @@ public class MongoTemplateTests {
 		@Id MyId id;
 	}
 
-	public static class Sample {
+	static class Sample {
 
 		@Id String id;
 		String field;