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.

Author: christophstrobl

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 @@ DBRef createDbRef(org.springframework.data.mongodb.core.mapping.DBRef annotation
 	 * @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.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.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.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 DBObject fetch(DBRef dbRef) {
 		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 @@ private synchronized Object resolve() {
 			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 @@ private Object readCollectionOrArray(TypeInformation<?> targetType, BasicDBList
 		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 @@ private Object readCollectionOrArray(TypeInformation<?> targetType, BasicDBList
 		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 @@ private <T> T potentiallyReadOrResolveDbRef(DBRef dbref, TypeInformation<?> type
 		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();
+	}
+
+	@SuppressWarnings("unchecked")
+	private <T> List<T> bulkReadAndConvertDBRefs(List<DBRef> dbrefs, TypeInformation<?> type, ObjectPath path,
+			final Class<?> rawType) {
 
-		if (readRef != null) {
-			maybeEmitEvent(new AfterLoadEvent<T>(readRef, (Class<T>) rawType, collectionName));
+		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 @@ DBObject readRef(DBRef ref) {
 		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.WriteConcern;
 import com.mongodb.WriteResult;
 
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.NoArgsConstructor;
+
 /**
  * Integration test for {@link MongoTemplate}.
  *
@@ -3368,6 +3372,50 @@ public void streamExecutionUsesExplicitCollectionName() {
 		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 boolean equals(Object obj) {
 
 	}
 
+	@Data
 	static class DocumentWithDBRefCollection {
 
 		@Id public String id;
@@ -3437,6 +3486,10 @@ static class DocumentWithDBRefCollection {
 
 		@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 @@ static class TypeWithMyId {
 		@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 @@ static class WithGeoJson {
 		String description;
 		GeoJsonPoint point;
 	}
-
 }