DATAMONGO-970 - MongoTemplate.remove(…) now correctly builds query for DBObjects.

If a DBObject was handed into MongoTemplate.remove(…) we previously failed to look up the id value to create a by-id-query. This commit adds explicit handling of DBObjects by looking up their _id field to obtain the id value.

Author: odrotbohm

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

@@ -1070,17 +1070,22 @@ public WriteResult remove(Object object, String collection) {
 	}
 
 	/**
-	 * Returns {@link Entry} containing the {@link MongoPersistentProperty} defining the {@literal id} as
-	 * {@link Entry#getKey()} and the {@link Id}s property value as its {@link Entry#getValue()}.
+	 * Returns {@link Entry} containing the field name of the id property as {@link Entry#getKey()} and the {@link Id}s
+	 * property value as its {@link Entry#getValue()}.
 	 * 
 	 * @param object
 	 * @return
 	 */
-	private Map.Entry<MongoPersistentProperty, Object> extractIdPropertyAndValue(Object object) {
+	private Entry<String, Object> extractIdPropertyAndValue(Object object) {
 
 		Assert.notNull(object, "Id cannot be extracted from 'null'.");
 
 		Class<?> objectType = object.getClass();
+
+		if (object instanceof DBObject) {
+			return Collections.singletonMap(ID_FIELD, ((DBObject) object).get(ID_FIELD)).entrySet().iterator().next();
+		}
+
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(objectType);
 		MongoPersistentProperty idProp = entity == null ? null : entity.getIdProperty();
 
@@ -1090,7 +1095,7 @@ private Map.Entry<MongoPersistentProperty, Object> extractIdPropertyAndValue(Obj
 
 		Object idValue = BeanWrapper.create(object, mongoConverter.getConversionService())
 				.getProperty(idProp, Object.class);
-		return Collections.singletonMap(idProp, idValue).entrySet().iterator().next();
+		return Collections.singletonMap(idProp.getFieldName(), idValue).entrySet().iterator().next();
 	}
 
 	/**
@@ -1101,8 +1106,8 @@ private Map.Entry<MongoPersistentProperty, Object> extractIdPropertyAndValue(Obj
 	 */
 	private Query getIdQueryFor(Object object) {
 
-		Map.Entry<MongoPersistentProperty, Object> id = extractIdPropertyAndValue(object);
-		return new Query(where(id.getKey().getFieldName()).is(id.getValue()));
+		Entry<String, Object> id = extractIdPropertyAndValue(object);
+		return new Query(where(id.getKey()).is(id.getValue()));
 	}
 
 	/**
@@ -1116,7 +1121,7 @@ private Query getIdInQueryFor(Collection<?> objects) {
 		Assert.notEmpty(objects, "Cannot create Query for empty collection.");
 
 		Iterator<?> it = objects.iterator();
-		Map.Entry<MongoPersistentProperty, Object> firstEntry = extractIdPropertyAndValue(it.next());
+		Entry<String, Object> firstEntry = extractIdPropertyAndValue(it.next());
 
 		ArrayList<Object> ids = new ArrayList<Object>(objects.size());
 		ids.add(firstEntry.getValue());
@@ -1125,7 +1130,7 @@ private Query getIdInQueryFor(Collection<?> objects) {
 			ids.add(extractIdPropertyAndValue(it.next()).getValue());
 		}
 
-		return new Query(where(firstEntry.getKey().getFieldName()).in(ids));
+		return new Query(where(firstEntry.getKey()).in(ids));
 	}
 
 	private void assertUpdateableIdIfNotSet(Object entity) {

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

@@ -2693,6 +2693,22 @@ public void shouldReuseExistingDBRefInQueryFromDbRefAssociationAfterLoad() {
 		assertThat(result.getContent().getName(), is(content.getName()));
 	}
 
+	/**
+	 * @see DATAMONGO-970
+	 */
+	@Test
+	public void insertsAndRemovesBasicDbObjectCorrectly() {
+
+		BasicDBObject object = new BasicDBObject("key", "value");
+		template.insert(object, "collection");
+
+		assertThat(object.get("_id"), is(notNullValue()));
+		assertThat(template.findAll(DBObject.class, "collection"), hasSize(1));
+
+		template.remove(object, "collection");
+		assertThat(template.findAll(DBObject.class, "collection"), hasSize(0));
+	}
+
 	static class DoucmentWithNamedIdField {
 
 		@Id String someIdKey;