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.

Author: odrotbohm

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

@@ -23,9 +23,11 @@
 import java.io.Serializable;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -142,8 +144,8 @@ public List<Document> bulkFetch(List<DBRef> refs) {
 		}
 
 		String collection = refs.iterator().next().getCollectionName();
+		List<Object> ids = new ArrayList<>(refs.size());
 
-		List<Object> ids = new ArrayList<Object>(refs.size());
 		for (DBRef ref : refs) {
 
 			if (!collection.equals(ref.getCollectionName())) {
@@ -155,10 +157,14 @@ public List<Document> bulkFetch(List<DBRef> refs) {
 		}
 
 		MongoDatabase db = mongoDbFactory.getDb();
-		List<Document> result = new ArrayList<>();
-		db.getCollection(collection).find(new Document("_id", new Document("$in", ids))).into(result);
-		result.sort(new DbRefByReferencePositionComparator(ids));
-		return result;
+
+		List<Document> result = db.getCollection(collection) //
+				.find(new Document("_id", new Document("$in", ids))) //
+				.into(new ArrayList<>());
+
+		return ids.stream() //
+				.flatMap(id -> documentWithId(id, result)) //
+				.collect(Collectors.toList());
 	}
 
 	/**
@@ -223,6 +229,20 @@ private boolean isLazyDbRef(MongoPersistentProperty property) {
 		return property.getDBRef() != null && property.getDBRef().lazy();
 	}
 
+	/**
+	 * Returns document with the given identifier from the given list of {@link Document}s.
+	 * 
+	 * @param identifier
+	 * @param documents
+	 * @return
+	 */
+	private static Stream<Document> documentWithId(Object identifier, Collection<Document> documents) {
+
+		return documents.stream() //
+				.filter(it -> it.get("_id").equals(identifier)) //
+				.limit(1);
+	}
+
 	/**
 	 * A {@link MethodInterceptor} that is used within a lazy loading proxy. The property resolving is delegated to a
 	 * {@link DbRefResolverCallback}. The resolving process is triggered by a method invocation on the proxy and is
@@ -449,37 +469,4 @@ private synchronized Object resolve() {
 			return result;
 		}
 	}
-
-	/**
-	 * {@link Comparator} for sorting {@link Document} that have been loaded in random order by a predefined list of
-	 * reference identifiers.
-	 *
-	 * @author Christoph Strobl
-	 * @author Oliver Gierke
-	 * @since 1.10
-	 */
-	private static class DbRefByReferencePositionComparator implements Comparator<Document> {
-
-		private final List<Object> reference;
-
-		/**
-		 * Creates a new {@link DbRefByReferencePositionComparator} for the given list of reference identifiers.
-		 *
-		 * @param referenceIds must not be {@literal null}.
-		 */
-		public DbRefByReferencePositionComparator(List<Object> referenceIds) {
-
-			Assert.notNull(referenceIds, "Reference identifiers must not be null!");
-			this.reference = new ArrayList<Object>(referenceIds);
-		}
-
-		/*
-		 * (non-Javadoc)
-		 * @see java.util.Comparator#compare(java.lang.Object, java.lang.Object)
-		 */
-		@Override
-		public int compare(Document o1, Document o2) {
-			return Integer.compare(reference.indexOf(o1.get("_id")), reference.indexOf(o2.get("_id")));
-		}
-	}
 }

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

@@ -15,18 +15,18 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.Matchers.*;
 import static org.hamcrest.Matchers.contains;
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertThat;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
 
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 
-import com.mongodb.client.FindIterable;
-import com.mongodb.client.MongoCollection;
-import com.mongodb.client.MongoDatabase;
 import org.bson.Document;
 import org.bson.types.ObjectId;
 import org.junit.Before;
@@ -43,6 +43,9 @@
 import org.springframework.data.mongodb.core.DocumentTestUtils;
 
 import com.mongodb.DBRef;
+import com.mongodb.client.FindIterable;
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.MongoDatabase;
 
 /**
  * Unit tests for {@link DefaultDbRefResolver}.
@@ -100,7 +103,7 @@ public void bulkFetchShouldThrowExceptionWhenUsingDifferntCollectionsWithinSetOf
 	@Test // DATAMONGO-1194
 	public void bulkFetchShouldReturnEarlyForEmptyLists() {
 
-		resolver.bulkFetch(Collections.<DBRef>emptyList());
+		resolver.bulkFetch(Collections.<DBRef> emptyList());
 
 		verify(collectionMock, never()).find(Mockito.any(Document.class));
 	}
@@ -115,6 +118,7 @@ public void bulkFetchShouldRestoreOriginalOrder() {
 		DBRef ref2 = new DBRef("collection-1", o2.get("_id"));
 
 		when(cursorMock.into(any())).then(new Answer<Object>() {
+
 			@Override
 			public Object answer(InvocationOnMock invocation) throws Throwable {
 
@@ -127,4 +131,17 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
 
 		assertThat(resolver.bulkFetch(Arrays.asList(ref1, ref2)), contains(o1, o2));
 	}
+
+	@Test // DATAMONGO-1765
+	public void bulkFetchContainsDuplicates() {
+
+		Document document = new Document("_id", new ObjectId());
+
+		DBRef ref1 = new DBRef("collection-1", document.get("_id"));
+		DBRef ref2 = new DBRef("collection-1", document.get("_id"));
+
+		when(cursorMock.into(any())).then(invocation -> Arrays.asList(document));
+
+		assertThat(resolver.bulkFetch(Arrays.asList(ref1, ref2))).containsExactly(document, document);
+	}
 }