DATAMONGO-511 - QueryMapper now maps associations correctly.

Complete overhaul of the QueryMapper to better handle complex scenarios like property paths and association references.

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

@@ -17,7 +17,6 @@ package org.springframework.data.mongodb.core.convert;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
 
 import org.bson.types.BasicBSONList;
@@ -76,47 +75,121 @@ public class QueryMapper {
 	 */
 	public DBObject getMappedObject(DBObject query, MongoPersistentEntity<?> entity) {
 
-		DBObject newDbo = new BasicDBObject();
+		if (isKeyWord(query)) {
+			return getMappedKeyword(query, entity);
+		}
+
+		DBObject result = new BasicDBObject();
 
 		for (String key : query.keySet()) {
 
+			MongoPersistentProperty targetProperty = getTargetProperty(key, entity);
 			String newKey = determineKey(key, entity);
 			Object value = query.get(key);
 
-			if (isIdKey(key, entity)) {
-				if (value instanceof DBObject) {
-					DBObject valueDbo = (DBObject) value;
-					if (valueDbo.containsField("$in") || valueDbo.containsField("$nin")) {
-						String inKey = valueDbo.containsField("$in") ? "$in" : "$nin";
-						List<Object> ids = new ArrayList<Object>();
-						for (Object id : (Iterable<?>) valueDbo.get(inKey)) {
-							ids.add(convertId(id));
-						}
-						valueDbo.put(inKey, ids.toArray(new Object[ids.size()]));
-					} else if (valueDbo.containsField("$ne")) {
-						valueDbo.put("$ne", convertId(valueDbo.get("$ne")));
-					} else {
-						value = getMappedObject((DBObject) value, null);
-					}
-				} else {
-					value = convertId(value);
-				}
-				newKey = "_id";
-			} else if (key.matches(N_OR_PATTERN)) {
-				// $or/$nor
-				Iterable<?> conditions = (Iterable<?>) value;
-				BasicBSONList newConditions = new BasicBSONList();
-				Iterator<?> iter = conditions.iterator();
-				while (iter.hasNext()) {
-					newConditions.add(getMappedObject((DBObject) iter.next(), entity));
-				}
-				value = newConditions;
-			}
-
-			newDbo.put(newKey, convertSimpleOrDBObject(value, null));
+			result.put(newKey, getMappedValue(value, targetProperty, newKey));
 		}
 
-		return newDbo;
+		return result;
+	}
+
+	/**
+	 * Returns the given {@link DBObject} representing a keyword by mapping the keyword's value.
+	 * 
+	 * @param query the {@link DBObject} representing a keyword (e.g. {@code $ne : … } )
+	 * @param entity
+	 * @return
+	 */
+	private DBObject getMappedKeyword(DBObject query, MongoPersistentEntity<?> entity) {
+
+		String newKey = query.keySet().iterator().next();
+		Object value = query.get(newKey);
+
+		// $or/$nor
+		if (newKey.matches(N_OR_PATTERN)) {
+
+			Iterable<?> conditions = (Iterable<?>) value;
+			BasicDBList newConditions = new BasicDBList();
+
+			for (Object condition : conditions) {
+				newConditions.add(getMappedObject((DBObject) condition, entity));
+			}
+
+			return new BasicDBObject(newKey, newConditions);
+		}
+
+		return new BasicDBObject(newKey, convertSimpleOrDBObject(value, entity));
+	}
+
+	/**
+	 * Returns the mapped value for the given source object assuming it's a value for the given
+	 * {@link MongoPersistentProperty}.
+	 * 
+	 * @param source the source object to be mapped
+	 * @param property the property the value is a value for
+	 * @param newKey the key the value will be bound to eventually
+	 * @return
+	 */
+	private Object getMappedValue(Object source, MongoPersistentProperty property, String newKey) {
+
+		if (property == null) {
+			return convertSimpleOrDBObject(source, null);
+		}
+
+		if (property.isIdProperty() || "_id".equals(newKey)) {
+
+			if (source instanceof DBObject) {
+				DBObject valueDbo = (DBObject) source;
+				if (valueDbo.containsField("$in") || valueDbo.containsField("$nin")) {
+					String inKey = valueDbo.containsField("$in") ? "$in" : "$nin";
+					List<Object> ids = new ArrayList<Object>();
+					for (Object id : (Iterable<?>) valueDbo.get(inKey)) {
+						ids.add(convertId(id));
+					}
+					valueDbo.put(inKey, ids.toArray(new Object[ids.size()]));
+				} else if (valueDbo.containsField("$ne")) {
+					valueDbo.put("$ne", convertId(valueDbo.get("$ne")));
+				} else {
+					return getMappedObject((DBObject) source, null);
+				}
+
+				return valueDbo;
+
+			} else {
+				return convertId(source);
+			}
+		}
+
+		if (property.isAssociation()) {
+			return isKeyWord(source) ? getMappedValue(getKeywordValue(source), property, newKey) : convertAssociation(source,
+					property);
+		}
+
+		return convertSimpleOrDBObject(source, mappingContext.getPersistentEntity(property));
+	}
+
+	private MongoPersistentProperty getTargetProperty(String key, MongoPersistentEntity<?> entity) {
+
+		if (isIdKey(key, entity)) {
+			return entity.getIdProperty();
+		}
+
+		PersistentPropertyPath<MongoPersistentProperty> path = getPath(key, entity);
+		return path == null ? null : path.getLeafProperty();
+	}
+
+	private PersistentPropertyPath<MongoPersistentProperty> getPath(String key, MongoPersistentEntity<?> entity) {
+
+		if (entity == null) {
+			return null;
+		}
+
+		try {
+			PropertyPath path = PropertyPath.from(key, entity.getTypeInformation());
+			return mappingContext.getPersistentPropertyPath(path);
+		} catch (PropertyReferenceException e) {
+			return null;
+		}
 	}
 
 	/**
@@ -128,17 +201,12 @@ public class QueryMapper {
 	 */
 	private String determineKey(String key, MongoPersistentEntity<?> entity) {
 
-		if (entity == null) {
-			return key;
+		if (entity == null && DEFAULT_ID_NAMES.contains(key)) {
+			return "_id";
 		}
 
-		try {
-			PropertyPath path = PropertyPath.from(key, entity.getTypeInformation());
-			PersistentPropertyPath<MongoPersistentProperty> propertyPath = mappingContext.getPersistentPropertyPath(path);
-			return propertyPath.toDotPath(MongoPersistentProperty.PropertyToFieldNameConverter.INSTANCE);
-		} catch (PropertyReferenceException e) {
-			return key;
-		}
+		PersistentPropertyPath<MongoPersistentProperty> path = getPath(key, entity);
+		return path == null ? key : path.toDotPath(MongoPersistentProperty.PropertyToFieldNameConverter.INSTANCE);
 	}
 
 	/**
@@ -161,6 +229,30 @@ public class QueryMapper {
 		return converter.convertToMongoType(source);
 	}
 
+	/**
+	 * Converts the given source assuming it's actually an association to anoter object.
+	 * 
+	 * @param source
+	 * @param property
+	 * @return
+	 */
+	private Object convertAssociation(Object source, MongoPersistentProperty property) {
+
+		if (property == null || !property.isAssociation()) {
+			return source;
+		}
+
+		if (source instanceof Iterable) {
+			BasicBSONList result = new BasicBSONList();
+			for (Object element : (Iterable<?>) source) {
+				result.add(converter.toDBRef(element, property));
+			}
+			return result;
+		}
+
+		return converter.toDBRef(source, property);
+	}
+
 	/**
 	 * Returns whether the given key will be considered an id key.
 	 * 
@@ -183,6 +275,34 @@ public class QueryMapper {
 		return DEFAULT_ID_NAMES.contains(key);
 	}
 
+	/**
+	 * Returns whether the given value is representing a query keyword.
+	 * 
+	 * @param value
+	 * @return
+	 */
+	private static boolean isKeyWord(Object value) {
+
+		if (!(value instanceof DBObject) || value instanceof BasicDBList) {
+			return false;
+		}
+
+		DBObject dbObject = (DBObject) value;
+		return dbObject.keySet().size() == 1 && dbObject.keySet().iterator().next().startsWith("$");
+	}
+
+	/**
+	 * Returns the value of the given source assuming it's a query keyword.
+	 * 
+	 * @param source
+	 * @return
+	 */
+	private static Object getKeywordValue(Object source) {
+
+		DBObject dbObject = (DBObject) source;
+		return dbObject.get(dbObject.keySet().iterator().next());
+	}
+
 	/**
 	 * Converts the given raw id value into either {@link ObjectId} or {@link String}.
 	 * 

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

@@ -35,6 +35,7 @@ import org.mockito.runners.MockitoJUnitRunner;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.mongodb.MongoDbFactory;
 import org.springframework.data.mongodb.core.Person;
+import org.springframework.data.mongodb.core.mapping.DBRef;
 import org.springframework.data.mongodb.core.mapping.Field;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
@@ -57,6 +58,7 @@ public class QueryMapperUnitTests {
 
 	QueryMapper mapper;
 	MongoMappingContext context;
+	MappingMongoConverter converter;
 
 	@Mock
 	MongoDbFactory factory;
@@ -66,7 +68,7 @@ public class QueryMapperUnitTests {
 
 		context = new MongoMappingContext();
 
-		MappingMongoConverter converter = new MappingMongoConverter(factory, context);
+		converter = new MappingMongoConverter(factory, context);
 		converter.afterPropertiesSet();
 
 		mapper = new QueryMapper(converter);
@@ -203,7 +205,7 @@ public class QueryMapperUnitTests {
 	}
 
 	@Test
-	public void doesNotHandleNestedFieldsWithDefaultIdNames() {
+	public void doesHandleNestedFieldsWithDefaultIdNames() {
 
 		BasicDBObject dbObject = new BasicDBObject("id", new ObjectId().toString());
 		dbObject.put("nested", new BasicDBObject("id", new ObjectId().toString()));
@@ -212,7 +214,7 @@ public class QueryMapperUnitTests {
 
 		DBObject result = mapper.getMappedObject(dbObject, entity);
 		assertThat(result.get("_id"), is(instanceOf(ObjectId.class)));
-		assertThat(((DBObject) result.get("nested")).get("id"), is(instanceOf(String.class)));
+		assertThat(((DBObject) result.get("nested")).get("_id"), is(instanceOf(ObjectId.class)));
 	}
 
 	/**
@@ -287,6 +289,35 @@ public class QueryMapperUnitTests {
 		assertThat(result.keySet().size(), is(1));
 	}
 
+	@Test
+	public void convertsAssociationCorrectly() {
+
+		Reference reference = new Reference();
+		reference.id = 5L;
+
+		Query query = query(where("reference").is(reference));
+		DBObject object = mapper.getMappedObject(query.getQueryObject(), context.getPersistentEntity(WithDBRef.class));
+
+		Object referenceObject = object.get("reference");
+
+		assertThat(referenceObject, is(instanceOf(com.mongodb.DBRef.class)));
+	}
+
+	@Test
+	public void convertsNestedAssociationCorrectly() {
+
+		Reference reference = new Reference();
+		reference.id = 5L;
+
+		Query query = query(where("withDbRef.reference").is(reference));
+		DBObject object = mapper.getMappedObject(query.getQueryObject(),
+				context.getPersistentEntity(WithDBRefWrapper.class));
+
+		Object referenceObject = object.get("withDbRef.reference");
+
+		assertThat(referenceObject, is(instanceOf(com.mongodb.DBRef.class)));
+	}
+
 	class IdWrapper {
 		Object id;
 	}
@@ -323,4 +354,20 @@ public class QueryMapperUnitTests {
 		@Field("foo")
 		CustomizedField field;
 	}
+
+	class WithDBRef {
+
+		@DBRef
+		Reference reference;
+	}
+
+	class Reference {
+
+		Long id;
+	}
+
+	class WithDBRefWrapper {
+
+		WithDBRef withDbRef;
+	}
 }