From 407087b3a74293a8890f48c42e46e1c2402e90b6 Mon Sep 17 00:00:00 2001
From: Oliver Gierke <ogierke@pivotal.io>
Date: Mon, 7 Aug 2017 17:32:10 +0200
Subject: [PATCH] DATAMONGO-1765 - DefaultDbRefResolver now maps duplicate
 references correctly.

On bulk resolution of a DBRef array we now map the resulting documents back to their ids to make sure that reoccurring identifiers are mapped to the corresponding documents.
---
 .../core/convert/DefaultDbRefResolver.java    | 40 ++++++++++++++++---
 .../DefaultDbRefResolverUnitTests.java        | 13 ++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolver.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolver.java
index 5f822ee28..deb49aafe 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolver.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolver.java
@@ -25,7 +25,9 @@ import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -46,7 +48,6 @@ import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
 
 import com.mongodb.BasicDBObject;
-import com.mongodb.BasicDBObjectBuilder;
 import com.mongodb.DB;
 import com.mongodb.DBObject;
 import com.mongodb.DBRef;
@@ -62,6 +63,8 @@ import com.mongodb.DBRef;
  */
 public class DefaultDbRefResolver implements DbRefResolver {
 
+	private static final String ID = "_id";
+
 	private final MongoDbFactory mongoDbFactory;
 	private final PersistenceExceptionTranslator exceptionTranslator;
 	private final ObjenesisStd objenesis;
@@ -144,10 +147,37 @@ public class DefaultDbRefResolver implements DbRefResolver {
 			ids.add(ref.getId());
 		}
 
+		Map<Object, DBObject> documentsById = getDocumentsById(ids, collection);
+		List<DBObject> result = new ArrayList<DBObject>(ids.size());
+
+		for (Object id : ids) {
+			result.add(documentsById.get(id));
+		}
+
+		return result;
+	}
+
+	/**
+	 * Returns all documents with the given ids contained in the given collection mapped by their ids.
+	 * 
+	 * @param ids must not be {@literal null}.
+	 * @param collection must not be {@literal null} or empty.
+	 * @return
+	 */
+	private Map<Object, DBObject> getDocumentsById(List<Object> ids, String collection) {
+
+		Assert.notNull(ids, "Ids must not be null!");
+		Assert.hasText(collection, "Collection must not be null or empty!");
+
 		DB db = mongoDbFactory.getDb();
-		List<DBObject> result = db.getCollection(collection)
-				.find(new BasicDBObjectBuilder().add("_id", new BasicDBObject("$in", ids)).get()).toArray();
-		Collections.sort(result, new DbRefByReferencePositionComparator(ids));
+		BasicDBObject query = new BasicDBObject(ID, new BasicDBObject("$in", ids));
+		List<DBObject> documents = db.getCollection(collection).find(query).toArray();
+		Map<Object, DBObject> result = new HashMap<Object, DBObject>(documents.size());
+
+		for (DBObject document : documents) {
+			result.put(document.get(ID), document);
+		}
+
 		return result;
 	}
 
@@ -469,7 +499,7 @@ public class DefaultDbRefResolver implements DbRefResolver {
 		 */
 		@Override
 		public int compare(DBObject o1, DBObject o2) {
-			return Integer.compare(reference.indexOf(o1.get("_id")), reference.indexOf(o2.get("_id")));
+			return Integer.compare(reference.indexOf(o1.get(ID)), reference.indexOf(o2.get(ID)));
 		}
 	}
 }
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolverUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolverUnitTests.java
index ba444fe25..c0becd2b7 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolverUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultDbRefResolverUnitTests.java
@@ -118,4 +118,17 @@ public class DefaultDbRefResolverUnitTests {
 
 		assertThat(resolver.bulkFetch(Arrays.asList(ref1, ref2)), contains(o1, o2));
 	}
+
+	@Test // DATAMONGO-1765
+	public void bulkFetchContainsDuplicates() {
+
+		DBObject document = new BasicDBObject("_id", new ObjectId());
+
+		DBRef ref1 = new DBRef("collection-1", document.get("_id"));
+		DBRef ref2 = new DBRef("collection-1", document.get("_id"));
+
+		when(cursorMock.toArray()).thenReturn(Arrays.asList(document));
+
+		assertThat(resolver.bulkFetch(Arrays.asList(ref1, ref2)), contains(document, document));
+	}
 }