DATAMONGO-1194 - Improve DBRef resolution for collections.

We now bulk load collections of referenced objects as long as they are stored in the same collection. This reduces database roundtrips and network traffic.

Original pull request: #377.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,8 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
+import java.util.List;
+
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 
@@ -64,4 +66,14 @@ public interface DbRefResolver {
 	 * @since 1.7
 	 */
 	DBObject fetch(DBRef dbRef);
+
+	/**
+	 * Loads a given {@link List} of {@link DBRef}s from the datasource in one batch. <br />
+	 * The {@link DBRef} elements in the list must not reference different collections.
+	 *
+	 * @param dbRefs must not be {@literal null}.
+	 * @return never {@literal null}.
+	 * @since 1.10
+	 */
+	List<DBObject> bulkFetch(List<DBRef> dbRefs);
 }

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

@@ -22,6 +22,10 @@ import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.Serializable;
 import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
 
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
@@ -31,6 +35,7 @@ import org.springframework.cglib.proxy.Enhancer;
 import org.springframework.cglib.proxy.Factory;
 import org.springframework.cglib.proxy.MethodProxy;
 import org.springframework.dao.DataAccessException;
+import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.dao.support.PersistenceExceptionTranslator;
 import org.springframework.data.mongodb.LazyLoadingException;
 import org.springframework.data.mongodb.MongoDbFactory;
@@ -40,6 +45,9 @@ import org.springframework.objenesis.ObjenesisStd;
 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;
 
@@ -109,6 +117,40 @@ public class DefaultDbRefResolver implements DbRefResolver {
 		return ReflectiveDBRefResolver.fetch(mongoDbFactory, dbRef);
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.convert.DbRefResolver#bulkFetch(java.util.List)
+	 */
+	@Override
+	public List<DBObject> bulkFetch(List<DBRef> refs) {
+
+		Assert.notNull(mongoDbFactory, "Factory must not be null!");
+		Assert.notNull(refs, "DBRef to fetch must not be null!");
+
+		if (refs.isEmpty()) {
+			return Collections.emptyList();
+		}
+
+		String collection = refs.iterator().next().getCollectionName();
+
+		List<Object> ids = new ArrayList<Object>(refs.size());
+		for (DBRef ref : refs) {
+
+			if (!collection.equals(ref.getCollectionName())) {
+				throw new InvalidDataAccessApiUsageException(
+						"DBRefs must all target the same collection for bulk fetch operation.");
+			}
+
+			ids.add(ref.getId());
+		}
+
+		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 DbRefByReferencePositionComperator(ids));
+		return result;
+	}
+
 	/**
 	 * Creates a proxy for the given {@link MongoPersistentProperty} using the given {@link DbRefResolverCallback} to
 	 * eventually resolve the value of the property.
@@ -395,4 +437,25 @@ public class DefaultDbRefResolver implements DbRefResolver {
 			return result;
 		}
 	}
+
+	/**
+	 * {@link Comparator} for sorting {@link DBObject} that have been loaded in random order by a predefined list of
+	 * reference ids.
+	 *
+	 * @author Christoph Strobl
+	 * @since 1.10
+	 */
+	private static class DbRefByReferencePositionComperator implements Comparator<DBObject> {
+
+		List<Object> reference;
+
+		public DbRefByReferencePositionComperator(List<Object> referenceIds) {
+			reference = new ArrayList<Object>(referenceIds);
+		}
+
+		@Override
+		public int compare(DBObject o1, DBObject o2) {
+			return Integer.compare(reference.indexOf(o1.get("_id")), reference.indexOf(o2.get("_id")));
+		}
+	}
 }

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

@@ -901,9 +901,14 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		Collection<Object> items = targetType.getType().isArray() ? new ArrayList<Object>()
 				: CollectionFactory.createCollection(collectionType, rawComponentType, sourceValue.size());
 
+		if (isCollectionOfDbRefWhereBulkFetchIsPossible(sourceValue) && !DBRef.class.equals(rawComponentType)) {
+			return bulkReadAndConvertDBRefs((List<DBRef>) (ArrayList) sourceValue, componentType, path, rawComponentType);
+		}
+
 		for (Object dbObjItem : sourceValue) {
 
 			if (dbObjItem instanceof DBRef) {
+
 				items.add(DBRef.class.equals(rawComponentType) ? dbObjItem
 						: readAndConvertDBRef((DBRef) dbObjItem, componentType, path, rawComponentType));
 			} else if (dbObjItem instanceof DBObject) {
@@ -916,6 +921,27 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		return getPotentiallyConvertedSimpleRead(items, targetType.getType());
 	}
 
+	private boolean isCollectionOfDbRefWhereBulkFetchIsPossible(Collection<Object> source) {
+
+		String collection = null;
+
+		for (Object dbObjItem : source) {
+
+			if (!(dbObjItem instanceof DBRef)) {
+				return false;
+			}
+
+			DBRef ref = (DBRef) dbObjItem;
+
+			if (collection != null && !collection.equals(ref.getCollectionName())) {
+				return false;
+			}
+			collection = ref.getCollectionName();
+		}
+
+		return true;
+	}
+
 	/**
 	 * Reads the given {@link DBObject} into a {@link Map}. will recursively resolve nested {@link Map}s as well.
 	 * 
@@ -1215,23 +1241,41 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		return (T) (object != null ? object : readAndConvertDBRef(dbref, type, path, rawType));
 	}
 
-	@SuppressWarnings("unchecked")
 	private <T> T readAndConvertDBRef(DBRef dbref, TypeInformation<?> type, ObjectPath path, final Class<?> rawType) {
 
-		final DBObject readRef = readRef(dbref);
-		final String collectionName = dbref.getCollectionName();
+		List<T> result = bulkReadAndConvertDBRefs(Collections.singletonList(dbref), type, path, rawType);
+		return CollectionUtils.isEmpty(result) ? null : result.iterator().next();
+	}
 
-		if (readRef != null) {
-			maybeEmitEvent(new AfterLoadEvent<T>(readRef, (Class<T>) rawType, collectionName));
+	@SuppressWarnings("unchecked")
+	private <T> List<T> bulkReadAndConvertDBRefs(List<DBRef> dbrefs, TypeInformation<?> type, ObjectPath path,
+			final Class<?> rawType) {
+
+		if (CollectionUtils.isEmpty(dbrefs)) {
+			return Collections.emptyList();
 		}
 
-		final T target = (T) read(type, readRef, path);
+		final List<DBObject> referencedRawDocuments = dbrefs.size() == 1
+				? Collections.singletonList(readRef(dbrefs.iterator().next())) : bulkReadRefs(dbrefs);
+		final String collectionName = dbrefs.iterator().next().getCollectionName();
 
-		if (target != null) {
-			maybeEmitEvent(new AfterConvertEvent<T>(readRef, target, collectionName));
+		List<T> targeList = new ArrayList<T>(dbrefs.size());
+
+		for (DBObject document : referencedRawDocuments) {
+
+			if (document != null) {
+				maybeEmitEvent(new AfterLoadEvent<T>(document, (Class<T>) rawType, collectionName));
+			}
+
+			final T target = (T) read(type, document, path);
+			targeList.add(target);
+
+			if (target != null) {
+				maybeEmitEvent(new AfterConvertEvent<T>(document, target, collectionName));
+			}
 		}
 
-		return target;
+		return targeList;
 	}
 
 	private void maybeEmitEvent(MongoMappingEvent<?> event) {
@@ -1255,6 +1299,17 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		return dbRefResolver.fetch(ref);
 	}
 
+	/**
+	 * Performs a bulk fetch operation for the given {@link DBRef}s.
+	 *
+	 * @param references must not be {@literal null}.
+	 * @return never {@literal null}.
+	 * @since 1.10
+	 */
+	List<DBObject> bulkReadRefs(List<DBRef> references) {
+		return dbRefResolver.bulkFetch(references);
+	}
+
 	/**
 	 * Marker class used to indicate we have a non root document object here that might be used within an update - so we
 	 * need to preserve type hints for potential nested elements but need to remove it on top level.

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

@@ -95,6 +95,10 @@ import com.mongodb.ReadPreference;
 import com.mongodb.WriteConcern;
 import com.mongodb.WriteResult;
 
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.NoArgsConstructor;
+
 /**
  * Integration test for {@link MongoTemplate}.
  *
@@ -3368,6 +3372,50 @@ public class MongoTemplateTests {
 		assertThat(stream.hasNext(), is(false));
 	}
 
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void shouldFetchListOfReferencesCorrectly() {
+
+		Sample one = new Sample("1", "jon snow");
+		Sample two = new Sample("2", "tyrion lannister");
+
+		template.save(one);
+		template.save(two);
+
+		DocumentWithDBRefCollection source = new DocumentWithDBRefCollection();
+		source.dbRefAnnotatedList = Arrays.asList(two, one);
+
+		template.save(source);
+
+		assertThat(template.findOne(query(where("id").is(source.id)), DocumentWithDBRefCollection.class), is(source));
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void shouldFetchListOfLazyReferencesCorrectly() {
+
+		Sample one = new Sample("1", "jon snow");
+		Sample two = new Sample("2", "tyrion lannister");
+
+		template.save(one);
+		template.save(two);
+
+		DocumentWithDBRefCollection source = new DocumentWithDBRefCollection();
+		source.lazyDbRefAnnotatedList = Arrays.asList(two, one);
+
+		template.save(source);
+
+		DocumentWithDBRefCollection target = template.findOne(query(where("id").is(source.id)),
+				DocumentWithDBRefCollection.class);
+
+		assertThat(target.lazyDbRefAnnotatedList, instanceOf(LazyLoadingProxy.class));
+		assertThat(target.getLazyDbRefAnnotatedList(), contains(two, one));
+	}
+
 	static class TypeWithNumbers {
 
 		@Id String id;
@@ -3427,6 +3475,7 @@ public class MongoTemplateTests {
 
 	}
 
+	@Data
 	static class DocumentWithDBRefCollection {
 
 		@Id public String id;
@@ -3437,6 +3486,10 @@ public class MongoTemplateTests {
 
 		@org.springframework.data.mongodb.core.mapping.DBRef //
 		public Sample dbRefProperty;
+
+		@Field("lazy_db_ref_list") /** @see DATAMONGO-1194 */
+		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) //
+		public List<Sample> lazyDbRefAnnotatedList;
 	}
 
 	static class DocumentWithCollection {
@@ -3528,13 +3581,13 @@ public class MongoTemplateTests {
 		@Id MyId id;
 	}
 
+	@EqualsAndHashCode
+	@NoArgsConstructor
 	static class Sample {
 
 		@Id String id;
 		String field;
 
-		public Sample() {}
-
 		public Sample(String id, String field) {
 			this.id = id;
 			this.field = field;
@@ -3729,5 +3782,4 @@ public class MongoTemplateTests {
 		String description;
 		GeoJsonPoint point;
 	}
-
 }

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

@@ -18,6 +18,7 @@ package org.springframework.data.mongodb.core.convert;
 import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
 import static org.mockito.Matchers.*;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.*;
 import static org.springframework.data.mongodb.core.convert.LazyLoadingTestUtils.*;
 
@@ -35,6 +36,7 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.runners.MockitoJUnitRunner;
 import org.springframework.data.annotation.AccessType;
 import org.springframework.data.annotation.AccessType.Type;
@@ -577,6 +579,68 @@ public class DbRefMappingMongoConverterUnitTests {
 		assertProxyIsResolved(result.dbRefToPlainObject, false);
 	}
 
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void shouldBulkFetchListOfReferences() {
+
+		String id1 = "1";
+		String id2 = "2";
+		String value = "val";
+
+		MappingMongoConverter converterSpy = spy(converter);
+		doReturn(Arrays.asList(new BasicDBObject("_id", id1).append("value", value),
+				new BasicDBObject("_id", id2).append("value", value))).when(converterSpy).bulkReadRefs(anyListOf(DBRef.class));
+
+		BasicDBObject dbo = new BasicDBObject();
+		ClassWithLazyDbRefs lazyDbRefs = new ClassWithLazyDbRefs();
+		lazyDbRefs.dbRefToConcreteCollection = new ArrayList<LazyDbRefTarget>(
+				Arrays.asList(new LazyDbRefTarget(id1, value), new LazyDbRefTarget(id2, value)));
+		converterSpy.write(lazyDbRefs, dbo);
+
+		ClassWithLazyDbRefs result = converterSpy.read(ClassWithLazyDbRefs.class, dbo);
+
+		assertProxyIsResolved(result.dbRefToConcreteCollection, false);
+		assertThat(result.dbRefToConcreteCollection.get(0).getId(), is(id1));
+		assertProxyIsResolved(result.dbRefToConcreteCollection, true);
+		assertThat(result.dbRefToConcreteCollection.get(1).getId(), is(id2));
+
+		verify(converterSpy, never()).readRef(Mockito.any(DBRef.class));
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void shouldFallbackToOneByOneFetchingWhenElementsInListOfReferencesPointToDifferentCollections() {
+
+		String id1 = "1";
+		String id2 = "2";
+		String value = "val";
+
+		MappingMongoConverter converterSpy = spy(converter);
+		doReturn(new BasicDBObject("_id", id1).append("value", value))
+				.doReturn(new BasicDBObject("_id", id2).append("value", value)).when(converterSpy)
+				.readRef(Mockito.any(DBRef.class));
+
+		BasicDBObject dbo = new BasicDBObject();
+		ClassWithLazyDbRefs lazyDbRefs = new ClassWithLazyDbRefs();
+		lazyDbRefs.dbRefToConcreteCollection = new ArrayList<LazyDbRefTarget>(
+				Arrays.asList(new LazyDbRefTarget(id1, value), new SerializableLazyDbRefTarget(id2, value)));
+		converterSpy.write(lazyDbRefs, dbo);
+
+		ClassWithLazyDbRefs result = converterSpy.read(ClassWithLazyDbRefs.class, dbo);
+
+		assertProxyIsResolved(result.dbRefToConcreteCollection, false);
+		assertThat(result.dbRefToConcreteCollection.get(0).getId(), is(id1));
+		assertProxyIsResolved(result.dbRefToConcreteCollection, true);
+		assertThat(result.dbRefToConcreteCollection.get(1).getId(), is(id2));
+
+		verify(converterSpy, times(2)).readRef(Mockito.any(DBRef.class));
+		verify(converterSpy, never()).bulkReadRefs(anyListOf(DBRef.class));
+	}
+
 	private Object transport(Object result) {
 		return SerializationUtils.deserialize(SerializationUtils.serialize(result));
 	}

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

@@ -0,0 +1,129 @@
+/*
+ * Copyright 2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.convert;
+
+import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
+import static org.hamcrest.collection.IsIterableWithSize.*;
+import static org.junit.Assert.*;
+import static org.mockito.Matchers.*;
+import static org.mockito.Mockito.*;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+import org.bson.types.ObjectId;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.dao.InvalidDataAccessApiUsageException;
+import org.springframework.data.mongodb.MongoDbFactory;
+import org.springframework.data.mongodb.core.DBObjectTestUtils;
+
+import com.mongodb.BasicDBObject;
+import com.mongodb.DB;
+import com.mongodb.DBCollection;
+import com.mongodb.DBCursor;
+import com.mongodb.DBObject;
+import com.mongodb.DBRef;
+
+/**
+ * @author Christoph Strobl
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class DefaultDbRefResolverUnitTests {
+
+	@Mock MongoDbFactory factoryMock;
+	@Mock DB dbMock;
+	@Mock DBCollection collectionMock;
+	@Mock DBCursor cursorMock;
+	DefaultDbRefResolver resolver;
+
+	@Before
+	public void setUp() {
+
+		when(factoryMock.getDb()).thenReturn(dbMock);
+		when(dbMock.getCollection(anyString())).thenReturn(collectionMock);
+		when(collectionMock.find(any(DBObject.class))).thenReturn(cursorMock);
+		when(cursorMock.toArray()).thenReturn(Collections.<DBObject> emptyList());
+
+		resolver = new DefaultDbRefResolver(factoryMock);
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	@SuppressWarnings("unchecked")
+	public void bulkFetchShouldLoadDbRefsCorrectly() {
+
+		DBRef ref1 = new DBRef("collection-1", new ObjectId());
+		DBRef ref2 = new DBRef("collection-1", new ObjectId());
+
+		resolver.bulkFetch(Arrays.asList(ref1, ref2));
+
+		ArgumentCaptor<DBObject> captor = ArgumentCaptor.forClass(DBObject.class);
+
+		verify(collectionMock, times(1)).find(captor.capture());
+
+		DBObject _id = DBObjectTestUtils.getAsDBObject(captor.getValue(), "_id");
+		Iterable<Object> $in = DBObjectTestUtils.getTypedValue(_id, "$in", Iterable.class);
+
+		assertThat($in, iterableWithSize(2));
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test(expected = InvalidDataAccessApiUsageException.class)
+	public void bulkFetchShouldThrowExceptionWhenUsingDifferntCollectionsWithinSetOfReferences() {
+
+		DBRef ref1 = new DBRef("collection-1", new ObjectId());
+		DBRef ref2 = new DBRef("collection-2", new ObjectId());
+
+		resolver.bulkFetch(Arrays.asList(ref1, ref2));
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void bulkFetchShouldReturnEarlyForEmptyLists() {
+
+		resolver.bulkFetch(Collections.<DBRef> emptyList());
+
+		verify(collectionMock, never()).find(any(DBObject.class));
+	}
+
+	/**
+	 * @see DATAMONGO-1194
+	 */
+	@Test
+	public void bulkFetchShouldRestoreOriginalOrder() {
+
+		DBObject o1 = new BasicDBObject("_id", new ObjectId());
+		DBObject o2 = new BasicDBObject("_id", new ObjectId());
+
+		DBRef ref1 = new DBRef("collection-1", o1.get("_id"));
+		DBRef ref2 = new DBRef("collection-1", o2.get("_id"));
+
+		when(cursorMock.toArray()).thenReturn(Arrays.asList(o2, o1));
+
+		assertThat(resolver.bulkFetch(Arrays.asList(ref1, ref2)), contains(o1, o2));
+	}
+}