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 @@
 
 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 QueryMapper(MongoConverter converter) {
  */
 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);
+result.put(newKey, getMappedValue(value, targetProperty, newKey));
+}
+
+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 {
-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));
+return getMappedObject((DBObject) source, null);
 }
-value = newConditions;
+
+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();
+}
 
-newDbo.put(newKey, convertSimpleOrDBObject(value, null));
+private PersistentPropertyPath<MongoPersistentProperty> getPath(String key, MongoPersistentEntity<?> entity) {
+
+if (entity == null) {
+return null;
 }
 
-return newDbo;
+try {
+PropertyPath path = PropertyPath.from(key, entity.getTypeInformation());
+return mappingContext.getPersistentPropertyPath(path);
+} catch (PropertyReferenceException e) {
+return null;
+}
 }
 
 /**
@@ -128,17 +201,12 @@ public DBObject getMappedObject(DBObject query, MongoPersistentEntity<?> entity)
  */
 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 @@ private Object convertSimpleOrDBObject(Object source, MongoPersistentEntity<?> e
 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 @@ private boolean isIdKey(String key, MongoPersistentEntity<?> entity) {
 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.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 void setUp() {
 
 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 void transformsArraysCorrectly() {
 }
 
 @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 void doesNotHandleNestedFieldsWithDefaultIdNames() {
 
 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 void returnsOriginalKeyIfNoPropertyReference() {
 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 @@ class CustomizedField {
 @Field("foo")
 CustomizedField field;
 }
+
+class WithDBRef {
+
+@DBRef
+Reference reference;
+}
+
+class Reference {
+
+Long id;
+}
+
+class WithDBRefWrapper {
+
+WithDBRef withDbRef;
+}
 }