DATAMONGO-1194 - Improve DBRef resolution for collections.

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-mongodb-parent</artifactId>
-	<version>1.10.0.BUILD-SNAPSHOT</version>
+	<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data MongoDB</name>

spring-data-mongodb-cross-store/pom.xml

@@ -6,7 +6,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>1.10.0.BUILD-SNAPSHOT</version>
+		<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 
@@ -48,7 +48,7 @@
 		<dependency>
 			<groupId>org.springframework.data</groupId>
 			<artifactId>spring-data-mongodb</artifactId>
-			<version>1.10.0.BUILD-SNAPSHOT</version>
+			<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 		</dependency>
 
 		<dependency>

spring-data-mongodb-distribution/pom.xml

@@ -13,7 +13,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>1.10.0.BUILD-SNAPSHOT</version>
+		<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb-log4j/pom.xml

@@ -5,7 +5,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>1.10.0.BUILD-SNAPSHOT</version>
+		<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-mongodb/pom.xml

@@ -11,7 +11,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-mongodb-parent</artifactId>
-		<version>1.10.0.BUILD-SNAPSHOT</version>
+		<version>1.10.0.DATAMONGO-1194-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

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,9 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
+import java.util.List;
+
+import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 
@@ -64,4 +67,16 @@ 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. The resulting {@link List} of
+	 * {@link DBObject} will reflect the ordering of the {@link DBRef} passed in.<br />
+	 * The {@link DBRef} elements in the list must not reference different collections.
+	 *
+	 * @param dbRefs must not be {@literal null}.
+	 * @return never {@literal null}.
+	 * @throws InvalidDataAccessApiUsageException in case not all {@link DBRef} target the same collection.
+	 * @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 DbRefByReferencePositionComparator(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 DbRefByReferencePositionComparator implements Comparator<DBObject> {
+
+		List<Object> reference;
+
+		public DbRefByReferencePositionComparator(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

@@ -883,6 +883,7 @@ public Object getValueInternal(MongoPersistentProperty prop, DBObject dbo, SpELE
 	 * @param path must not be {@literal null}.
 	 * @return the converted {@link Collection} or array, will never be {@literal null}.
 	 */
+	@SuppressWarnings({ "rawtypes", "unchecked" })
 	private Object readCollectionOrArray(TypeInformation<?> targetType, BasicDBList sourceValue, ObjectPath path) {
 
 		Assert.notNull(targetType, "Target type must not be null!");
@@ -901,9 +902,14 @@ private Object readCollectionOrArray(TypeInformation<?> targetType, BasicDBList
 		Collection<Object> items = targetType.getType().isArray() ? new ArrayList<Object>()
 				: CollectionFactory.createCollection(collectionType, rawComponentType, sourceValue.size());
 
+		if (!DBRef.class.equals(rawComponentType) && isCollectionOfDbRefWhereBulkFetchIsPossible(sourceValue)) {
+			return bulkReadAndConvertDBRefs((List<DBRef>) (List) (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) {
@@ -941,6 +947,11 @@ protected Map<Object, Object> readMap(TypeInformation<?> type, DBObject dbObject
 		Map<Object, Object> map = CollectionFactory.createMap(mapType, rawKeyType, dbObject.keySet().size());
 		Map<String, Object> sourceMap = dbObject.toMap();
 
+		if (!DBRef.class.equals(rawValueType) && isCollectionOfDbRefWhereBulkFetchIsPossible(sourceMap.values())) {
+			bulkReadAndConvertDBRefMapIntoTarget(valueType, rawValueType, sourceMap, map);
+			return map;
+		}
+
 		for (Entry<String, Object> entry : sourceMap.entrySet()) {
 			if (typeMapper.isTypeKey(entry.getKey())) {
 				continue;
@@ -1215,23 +1226,77 @@ 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", "rawtypes" })
+	private void bulkReadAndConvertDBRefMapIntoTarget(TypeInformation<?> valueType, Class<?> rawValueType,
+			Map<String, Object> sourceMap, Map<Object, Object> targetMap) {
+
+		LinkedHashMap<String, Object> referenceMap = new LinkedHashMap<String, Object>(sourceMap);
+		List<Object> convertedObjects = bulkReadAndConvertDBRefs((List<DBRef>) new ArrayList(referenceMap.values()),
+				valueType, ObjectPath.ROOT, rawValueType);
 
-		if (readRef != null) {
-			maybeEmitEvent(new AfterLoadEvent<T>(readRef, (Class<T>) rawType, collectionName));
+		int index = 0;
+		for (String key : referenceMap.keySet()) {
+			targetMap.put(key, convertedObjects.get(index));
+			index++;
 		}
+	}
 
-		final T target = (T) read(type, readRef, path);
+	@SuppressWarnings("unchecked")
+	private <T> List<T> bulkReadAndConvertDBRefs(List<DBRef> dbrefs, TypeInformation<?> type, ObjectPath path,
+			final Class<?> rawType) {
 
-		if (target != null) {
-			maybeEmitEvent(new AfterConvertEvent<T>(readRef, target, collectionName));
+		if (CollectionUtils.isEmpty(dbrefs)) {
+			return Collections.emptyList();
+		}
+
+		final List<DBObject> referencedRawDocuments = dbrefs.size() == 1
+				? Collections.singletonList(readRef(dbrefs.iterator().next())) : bulkReadRefs(dbrefs);
+		final String collectionName = dbrefs.iterator().next().getCollectionName();
+
+		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 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;
 	}
 
 	private void maybeEmitEvent(MongoMappingEvent<?> event) {
@@ -1255,6 +1320,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.