From 170081137a82d3e1d3e247aca6c86b6a0bec44f0 Mon Sep 17 00:00:00 2001
From: Oliver Gierke <info@olivergierke.de>
Date: Wed, 14 Dec 2011 20:31:08 +0100
Subject: [PATCH] DATAMONGO-346 - Fixed id-handling in queries.

In case a query referencing an entity's id needs massaging before being executed (e.g. Strings that can be ObjectID) the massaging failed in case the query was built using _id already as we did not detect that one being an id reference as we compared it to the entity's id property name. We now also compare against it's field name.
---
 .../data/mongodb/core/QueryMapper.java        | 40 ++++++++++++-------
 .../data/mongodb/core/MongoTemplateTests.java | 30 +++++++++++++-
 2 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryMapper.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryMapper.java
index 367e227c2..96fca71ef 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryMapper.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryMapper.java
@@ -16,6 +16,7 @@
 package org.springframework.data.mongodb.core;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
 
@@ -26,6 +27,7 @@ import org.springframework.core.convert.ConversionService;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.util.Assert;
 
 import com.mongodb.BasicDBObject;
@@ -62,22 +64,13 @@ public class QueryMapper {
 	 * @return
 	 */
 	public DBObject getMappedObject(DBObject query, MongoPersistentEntity<?> entity) {
-		
-		String idKey = null;
-		
-		if (null != entity && entity.getIdProperty() != null) {
-			idKey = entity.getIdProperty().getName();
-		} else if (query.containsField("id")) {
-			idKey = "id";
-		} else if (query.containsField("_id")) {
-			idKey = "_id";
-		}
 
 		DBObject newDbo = new BasicDBObject();
+
 		for (String key : query.keySet()) {
 			String newKey = key;
 			Object value = query.get(key);
-			if (key.equals(idKey)) {
+			if (isIdKey(key, entity)) {
 				if (value instanceof DBObject) {
 					DBObject valueDbo = (DBObject) value;
 					if (valueDbo.containsField("$in") || valueDbo.containsField("$nin")) {
@@ -106,13 +99,30 @@ public class QueryMapper {
 			} else if (key.equals("$ne")) {
 				value = convertId(value);
 			}
-			
+
 			newDbo.put(newKey, converter.convertToMongoType(value));
 		}
-		
+
 		return newDbo;
 	}
 
+	/**
+	 * Returns whether the given key will be considered an id key.
+	 * 
+	 * @param key
+	 * @param entity
+	 * @return
+	 */
+	private boolean isIdKey(String key, MongoPersistentEntity<?> entity) {
+
+		if (null != entity && entity.getIdProperty() != null) {
+			MongoPersistentProperty idProperty = entity.getIdProperty();
+			return idProperty.getName().equals(key) || idProperty.getFieldName().equals(key);
+		}
+
+		return Arrays.asList("id", "_id").contains(key);
+	}
+
 	/**
 	 * Converts the given raw id value into either {@link ObjectId} or {@link String}.
 	 * 
@@ -120,13 +130,13 @@ public class QueryMapper {
 	 * @return
 	 */
 	public Object convertId(Object id) {
-		
+
 		try {
 			return conversionService.convert(id, ObjectId.class);
 		} catch (ConversionException e) {
 			// Ignore
 		}
-		
+
 		return converter.convertToMongoType(id);
 	}
 }
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
index 79ef4ce7a..8bc06becd 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java
@@ -39,6 +39,7 @@ import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.dao.DataAccessException;
+import org.springframework.data.annotation.Id;
 import org.springframework.data.annotation.PersistenceConstructor;
 import org.springframework.data.mongodb.InvalidMongoDbApiUsageException;
 import org.springframework.data.mongodb.MongoDbFactory;
@@ -428,12 +429,12 @@ public class MongoTemplateTests {
 		assertThat(p.getAge(), is(26));
 		p = template.findOne(query, Person.class);
 		assertThat(p.getAge(), is(27));
-		
+
 		Query query2 = new Query(Criteria.where("firstName").is("Mary"));
 		p = template.findAndModify(query2, update, new FindAndModifyOptions().returnNew(true).upsert(true), Person.class);
 		assertThat(p.getFirstName(), is("Mary"));
 		assertThat(p.getAge(), is(1));
-		
+
 	}
 
 	@Test
@@ -1060,6 +1061,31 @@ public class MongoTemplateTests {
 		assertThat(template.findAll(Person.class, "mycollection").isEmpty(), is(true));
 	}
 
+	/**
+	 * @see DATADOC-349
+	 */
+	@Test
+	public void removesEntityWithAnnotatedIdIfIdNeedsMassaging() {
+
+		String id = new ObjectId().toString();
+
+		Sample sample = new Sample();
+		sample.id = id;
+
+		template.save(sample);
+
+		assertThat(template.findOne(query(where("id").is(id)), Sample.class).id, is(id));
+
+		template.remove(sample);
+		assertThat(template.findOne(query(where("id").is(id)), Sample.class), is(nullValue()));
+	}
+
+	public class Sample {
+
+		@Id
+		String id;
+	}
+
 	public class TestClass {
 
 		private DateTime myDate;