DATAMONGO-404 - Fixed Update.pull(…) handling to work with DBRefs.

We now support pointing to DBRef-mapped properties in Update.pull(…) and also allow to refer to the id of the DBRef to avoid having to create an instance of the entity.

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

@@ -26,6 +26,7 @@ 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.Association;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.PropertyPath;
 import org.springframework.data.mapping.PropertyReferenceException;
@@ -187,8 +188,8 @@ public class QueryMapper {
 		boolean needsAssociationConversion = property.isAssociation() && !keyword.isExists();
 		Object value = keyword.getValue();
 
-		Object convertedValue = needsAssociationConversion ? convertAssociation(value, property.getProperty())
-				: getMappedValue(property.with(keyword.getKey()), value);
+		Object convertedValue = needsAssociationConversion ? convertAssociation(value, property) : getMappedValue(
+				property.with(keyword.getKey()), value);
 
 		return new BasicDBObject(keyword.key, convertedValue);
 	}
@@ -235,7 +236,7 @@ public class QueryMapper {
 		}
 
 		if (isAssociationConversionNecessary(documentField, value)) {
-			return convertAssociation(value, documentField.getProperty());
+			return convertAssociation(value, documentField);
 		}
 
 		return convertSimpleOrDBObject(value, documentField.getPropertyEntity());
@@ -251,9 +252,10 @@ public class QueryMapper {
 	 * @param value
 	 * @return
 	 */
-	private boolean isAssociationConversionNecessary(Field documentField, Object value) {
+	protected boolean isAssociationConversionNecessary(Field documentField, Object value) {
 		return documentField.isAssociation() && value != null
-				&& documentField.getProperty().getActualType().isAssignableFrom(value.getClass());
+				&& (documentField.getProperty().getActualType().isAssignableFrom(value.getClass()) //
+				|| documentField.getPropertyEntity().getIdProperty().getActualType().isAssignableFrom(value.getClass()));
 	}
 
 	/**
@@ -288,6 +290,10 @@ public class QueryMapper {
 		return converter.convertToMongoType(source);
 	}
 
+	protected Object convertAssociation(Object source, Field field) {
+		return convertAssociation(source, field.getProperty());
+	}
+
 	/**
 	 * Converts the given source assuming it's actually an association to another object.
 	 * 
@@ -297,15 +303,14 @@ public class QueryMapper {
 	 */
 	private Object convertAssociation(Object source, MongoPersistentProperty property) {
 
-		if (property == null || !property.isAssociation() || source == null || source instanceof DBRef
-				|| !property.isEntity()) {
+		if (property == null || source == null || source instanceof DBRef) {
 			return source;
 		}
 
 		if (source instanceof Iterable) {
 			BasicDBList result = new BasicDBList();
 			for (Object element : (Iterable<?>) source) {
-				result.add(element instanceof DBRef ? element : converter.toDBRef(element, property));
+				result.add(createDbRefFor(element, property));
 			}
 			return result;
 		}
@@ -314,12 +319,20 @@ public class QueryMapper {
 			BasicDBObject result = new BasicDBObject();
 			DBObject dbObject = (DBObject) source;
 			for (String key : dbObject.keySet()) {
-				Object o = dbObject.get(key);
-				result.put(key, o instanceof DBRef ? o : converter.toDBRef(o, property));
+				result.put(key, createDbRefFor(dbObject.get(key), property));
 			}
 			return result;
 		}
 
+		return createDbRefFor(source, property);
+	}
+
+	private DBRef createDbRefFor(Object source, MongoPersistentProperty property) {
+
+		if (source instanceof DBRef) {
+			return (DBRef) source;
+		}
+
 		return converter.toDBRef(source, property);
 	}
 
@@ -502,6 +515,24 @@ public class QueryMapper {
 		public String getMappedKey() {
 			return isIdField() ? ID_KEY : name;
 		}
+
+		/**
+		 * Checks if property is part of database reference.
+		 * 
+		 * @return
+		 */
+		public boolean isPartOfAssociation() {
+			return false;
+		}
+
+		/**
+		 * Get {@link Association} pointing to property
+		 * 
+		 * @return null if not available
+		 */
+		public Association<MongoPersistentProperty> findAssociation() {
+			return null;
+		}
 	}
 
 	/**
@@ -594,6 +625,34 @@ public class QueryMapper {
 			return property == null ? false : property.isAssociation();
 		}
 
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.core.convert.QueryMapper.Field#isPartOfAssociation()
+		 */
+		@Override
+		public boolean isPartOfAssociation() {
+			return findAssociation() != null;
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.core.convert.QueryMapper.Field#findAssociation()
+		 */
+		@Override
+		public Association<MongoPersistentProperty> findAssociation() {
+			if (isAssociation()) {
+				return property.getAssociation();
+			}
+			if (this.path != null) {
+				for (MongoPersistentProperty p : this.path) {
+					if (p != null && p.isAssociation()) {
+						return p.getAssociation();
+					}
+				}
+			}
+			return null;
+		}
+
 		/*
 		 * (non-Javadoc)
 		 * @see org.springframework.data.mongodb.core.convert.QueryMapper.Field#getTargetKey()

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

@@ -21,6 +21,8 @@ import java.util.Iterator;
 import java.util.Map.Entry;
 
 import org.springframework.core.convert.converter.Converter;
+import org.springframework.data.mapping.Association;
+import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
@@ -32,6 +34,7 @@ import org.springframework.util.Assert;
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
+import com.mongodb.DBRef;
 
 /**
  * A subclass of {@link QueryMapper} that retains type information on the mongo types.
@@ -76,7 +79,7 @@ public class UpdateMapper extends QueryMapper {
 	protected Entry<String, Object> getMappedObjectForField(Field field, Object rawValue) {
 
 		if (!isUpdateModifier(rawValue)) {
-			return super.getMappedObjectForField(field, rawValue);
+			return super.getMappedObjectForField(field, getMappedValue(field, rawValue));
 		}
 
 		Object value = null;
@@ -102,6 +105,27 @@ public class UpdateMapper extends QueryMapper {
 		return Collections.singletonMap(field.getMappedKey(), value).entrySet().iterator().next();
 	}
 
+	@Override
+	protected Object convertAssociation(Object source, Field field) {
+
+		if (source instanceof DBRef || field.isAssociation() || !field.isPartOfAssociation()) {
+			return super.convertAssociation(source, field);
+		}
+
+		MongoPersistentProperty property = field.getProperty();
+		PersistentEntity<?, MongoPersistentProperty> owner = property.getOwner();
+		if (owner instanceof MongoPersistentEntity) {
+			return new DBRef(null, ((MongoPersistentEntity<?>) owner).getCollection(), source);
+		}
+
+		throw new IllegalArgumentException(String.format("Expected MongoPeristentEntity but found '%s'.", owner.getClass()));
+	}
+
+	@Override
+	protected boolean isAssociationConversionNecessary(Field documentField, Object value) {
+		return super.isAssociationConversionNecessary(documentField, value) || documentField.isPartOfAssociation();
+	}
+
 	private boolean isUpdateModifier(Object value) {
 		return value instanceof Modifier || value instanceof Modifiers;
 	}
@@ -159,9 +183,41 @@ public class UpdateMapper extends QueryMapper {
 		 */
 		@Override
 		protected Converter<MongoPersistentProperty, String> getPropertyConverter() {
+
+			if (this.isPartOfAssociation()) {
+				return new AssociationPropertyConverter(key, this.findAssociation());
+			}
 			return new UpdatePropertyConverter(key);
 		}
 
+		@Override
+		public String getMappedKey() {
+
+			if (isPartOfAssociation()) {
+				return getPropertyConverter().convert(null);
+			}
+			return super.getMappedKey();
+		}
+
+		private static class AssociationPropertyConverter implements Converter<MongoPersistentProperty, String> {
+
+			Association<MongoPersistentProperty> association;
+			String key;
+
+			public AssociationPropertyConverter(String key, Association<MongoPersistentProperty> field) {
+				this.key = key;
+				this.association = field;
+			}
+
+			@Override
+			public String convert(MongoPersistentProperty source) {
+				String fieldName = association.getInverse().getFieldName();
+				String newKey = key.substring(0, key.indexOf(fieldName) + fieldName.length());
+				return newKey;
+			}
+
+		}
+
 		/**
 		 * Special {@link Converter} for {@link MongoPersistentProperty} instances that will concatenate the {@literal $}
 		 * contained in the source update key.

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

@@ -2292,6 +2292,8 @@ public class MongoTemplateTests {
 	}
 
 	/**
+	 * <<<<<<< HEAD
+	 * 
 	 * @see DATAMONOGO-828
 	 */
 	@Test
@@ -2325,6 +2327,83 @@ public class MongoTemplateTests {
 		DocumentWithMultipleCollections result = template.findOne(findQuery, DocumentWithMultipleCollections.class);
 		assertThat(result.string1, hasItems("spring", "data", "mongodb"));
 		assertThat(result.string2, hasItems("one", "two", "three"));
+
+	}
+
+	/**
+	 * @see DATAMONGO-404
+	 */
+	@Test
+	public void updateWithPullShouldRemoveNestedItemFromDbRefAnnotatedCollection() {
+
+		template.dropCollection(DocumentWithDBRefCollection.class);
+
+		Sample sample1 = new Sample("1", "A");
+		Sample sample2 = new Sample("2", "B");
+		template.save(sample1);
+		template.save(sample2);
+
+		DocumentWithDBRefCollection doc = new DocumentWithDBRefCollection();
+		doc.id = "1";
+		doc.dbRefAnnotatedList = Arrays.asList( //
+				sample1, //
+				sample2 //
+				);
+		template.save(doc);
+
+		Update update = new Update().pull("dbRefAnnotatedList", doc.dbRefAnnotatedList.get(1));
+
+		Query qry = query(where("id").is("1"));
+		template.updateFirst(qry, update, DocumentWithDBRefCollection.class);
+
+		DocumentWithDBRefCollection result = template.findOne(qry, DocumentWithDBRefCollection.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.dbRefAnnotatedList, hasSize(1));
+		assertThat(result.dbRefAnnotatedList.get(0), is(notNullValue()));
+		assertThat(result.dbRefAnnotatedList.get(0).id, is((Object) "1"));
+	}
+
+	/**
+	 * @see DATAMONGO-404
+	 */
+	@Test
+	public void updateWithPullShouldRemoveNestedItemFromDbRefAnnotatedCollectionWhenGivenAnIdValueOfComponentTypeEntity() {
+
+		template.dropCollection(DocumentWithDBRefCollection.class);
+
+		Sample sample1 = new Sample("1", "A");
+		Sample sample2 = new Sample("2", "B");
+		template.save(sample1);
+		template.save(sample2);
+
+		DocumentWithDBRefCollection doc = new DocumentWithDBRefCollection();
+		doc.id = "1";
+		doc.dbRefAnnotatedList = Arrays.asList( //
+				sample1, //
+				sample2 //
+				);
+		template.save(doc);
+
+		Update update = new Update().pull("dbRefAnnotatedList.id", "2");
+
+		Query qry = query(where("id").is("1"));
+		template.updateFirst(qry, update, DocumentWithDBRefCollection.class);
+
+		DocumentWithDBRefCollection result = template.findOne(qry, DocumentWithDBRefCollection.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.dbRefAnnotatedList, hasSize(1));
+		assertThat(result.dbRefAnnotatedList.get(0), is(notNullValue()));
+		assertThat(result.dbRefAnnotatedList.get(0).id, is((Object) "1"));
+	}
+
+	static class DocumentWithDBRefCollection {
+
+		@Id public String id;
+
+		@org.springframework.data.mongodb.core.mapping.DBRef//
+		public List<Sample> dbRefAnnotatedList;
 	}
 
 	static class DocumentWithCollection {
@@ -2388,6 +2467,13 @@ public class MongoTemplateTests {
 
 		@Id String id;
 		String field;
+
+		public Sample() {}
+
+		public Sample(String id, String field) {
+			this.id = id;
+			this.field = field;
+		}
 	}
 
 	static class TestClass {