DATAMONGO-2004 - Support lazy DBRef resolution through constructor creation of the enclosing entity.

We now respect eager/lazy loading preferences of the annotated association property when the enclosing entity is created through its constructor and the reference is passed as constructor argument.

Previously, we eagerly resolved DBRefs and passed the resolved value to the constructor.

Original Pull Request: #571

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

@@ -58,6 +58,14 @@ class DocumentAccessor {
 		this.document = document;
 	}
 
+	/**
+	 * @return the underlying {@link Bson document}.
+	 * @since 2.1
+	 */
+	public Bson getDocument() {
+		return this.document;
+	}
+
 	/**
 	 * Puts the given value into the backing {@link Document} based on the coordinates defined through the given
 	 * {@link MongoPersistentProperty}. By default this will be the plain field name. But field names might also consist

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

@@ -255,7 +255,8 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 	private ParameterValueProvider<MongoPersistentProperty> getParameterProvider(MongoPersistentEntity<?> entity,
 			Bson source, DefaultSpELExpressionEvaluator evaluator, ObjectPath path) {
 
-		MongoDbPropertyValueProvider provider = new MongoDbPropertyValueProvider(source, evaluator, path);
+		AssociationAwareMongoDbPropertyValueProvider provider = new AssociationAwareMongoDbPropertyValueProvider(source,
+				evaluator, path);
 		PersistentEntityParameterValueProvider<MongoPersistentProperty> parameterProvider = new PersistentEntityParameterValueProvider<>(
 				entity, provider, path.getCurrentObject());
 
@@ -1273,9 +1274,9 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 	 */
 	class MongoDbPropertyValueProvider implements PropertyValueProvider<MongoPersistentProperty> {
 
-		private final DocumentAccessor source;
-		private final SpELExpressionEvaluator evaluator;
-		private final ObjectPath path;
+		final DocumentAccessor source;
+		final SpELExpressionEvaluator evaluator;
+		final ObjectPath path;
 
 		/**
 		 * Creates a new {@link MongoDbPropertyValueProvider} for the given source, {@link SpELExpressionEvaluator} and
@@ -1320,6 +1321,7 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		 * @see org.springframework.data.convert.PropertyValueProvider#getPropertyValue(org.springframework.data.mapping.PersistentProperty)
 		 */
 		@Nullable
+		@SuppressWarnings("unchecked")
 		public <T> T getPropertyValue(MongoPersistentProperty property) {
 
 			String expression = property.getSpelExpression();
@@ -1333,6 +1335,50 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		}
 	}
 
+	/**
+	 * {@link PropertyValueProvider} that is aware of {@link MongoPersistentProperty#isAssociation()} and that delegates
+	 * resolution to {@link DbRefResolver}.
+	 *
+	 * @author Mark Paluch
+	 * @since 2.1
+	 */
+	class AssociationAwareMongoDbPropertyValueProvider extends MongoDbPropertyValueProvider {
+
+		/**
+		 * Creates a new {@link AssociationAwareMongoDbPropertyValueProvider} for the given source,
+		 * {@link SpELExpressionEvaluator} and {@link ObjectPath}.
+		 *
+		 * @param source must not be {@literal null}.
+		 * @param evaluator must not be {@literal null}.
+		 * @param path must not be {@literal null}.
+		 */
+		public AssociationAwareMongoDbPropertyValueProvider(Bson source, SpELExpressionEvaluator evaluator,
+				ObjectPath path) {
+			super(source, evaluator, path);
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.convert.PropertyValueProvider#getPropertyValue(org.springframework.data.mapping.PersistentProperty)
+		 */
+		@Nullable
+		@SuppressWarnings("unchecked")
+		public <T> T getPropertyValue(MongoPersistentProperty property) {
+
+			T value = super.getPropertyValue(property);
+
+			if (value == null || !property.isAssociation()) {
+				return value;
+			}
+
+			DbRefResolverCallback callback = new DefaultDbRefResolverCallback(source.getDocument(), path, evaluator,
+					MappingMongoConverter.this);
+			DBRef dbref = value instanceof DBRef ? (DBRef) value : null;
+
+			return (T) dbRefResolver.resolveDbRef(property, dbref, callback, dbRefProxyHandler);
+		}
+	}
+
 	/**
 	 * Extension of {@link SpELExpressionParameterValueProvider} to recursively trigger value conversion on the raw
 	 * resolved SpEL value.

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

@@ -3069,8 +3069,8 @@ public class MongoTemplateTests {
 		assertThat(contentLoaded.dbrefMessage.id, is(messageLoaded.id));
 	}
 
-	@Test // DATAMONGO-1287
-	public void shouldReuseAlreadyResolvedLazyLoadedDBRefWhenUsedAsPersistenceConstrcutorArgument() {
+	@Test // DATAMONGO-1287, DATAMONGO-2004
+	public void shouldReuseAlreadyResolvedLazyLoadedDBRefWhenUsedAsPersistenceConstructorArgument() {
 
 		Document docInCtor = new Document();
 		docInCtor.id = "doc-in-ctor";
@@ -3083,7 +3083,7 @@ public class MongoTemplateTests {
 
 		DocumentWithLazyDBrefUsedInPresistenceConstructor loaded = template.findOne(query(where("id").is(source.id)),
 				DocumentWithLazyDBrefUsedInPresistenceConstructor.class);
-		assertThat(loaded.refToDocUsedInCtor, not(instanceOf(LazyLoadingProxy.class)));
+		assertThat(loaded.refToDocUsedInCtor, instanceOf(LazyLoadingProxy.class));
 		assertThat(loaded.refToDocNotUsedInCtor, nullValue());
 	}
 
@@ -3106,8 +3106,8 @@ public class MongoTemplateTests {
 		assertThat(loaded.refToDocUsedInCtor, nullValue());
 	}
 
-	@Test // DATAMONGO-1287
-	public void shouldRespectParamterValueWhenAttemptingToReuseLazyLoadedDBRefUsedInPersistenceConstrcutor() {
+	@Test // DATAMONGO-1287, DATAMONGO-2004
+	public void shouldRespectParameterValueWhenAttemptingToReuseLazyLoadedDBRefUsedInPersistenceConstructor() {
 
 		Document docInCtor = new Document();
 		docInCtor.id = "doc-in-ctor";
@@ -3125,7 +3125,7 @@ public class MongoTemplateTests {
 
 		DocumentWithLazyDBrefUsedInPresistenceConstructor loaded = template.findOne(query(where("id").is(source.id)),
 				DocumentWithLazyDBrefUsedInPresistenceConstructor.class);
-		assertThat(loaded.refToDocUsedInCtor, not(instanceOf(LazyLoadingProxy.class)));
+		assertThat(loaded.refToDocUsedInCtor, instanceOf(LazyLoadingProxy.class));
 		assertThat(loaded.refToDocNotUsedInCtor, instanceOf(LazyLoadingProxy.class));
 	}
 
@@ -3384,6 +3384,73 @@ public class MongoTemplateTests {
 		assertThat(target.lazyDbRefAnnotatedMap.values(), contains(two, one));
 	}
 
+	@Test // DATAMONGO-2004
+	public void shouldFetchLazyReferenceWithConstructorCreationCorrectly() {
+
+		Sample one = new Sample("1", "jon snow");
+
+		template.save(one);
+
+		DocumentWithLazyDBRefsAndConstructorCreation source = new DocumentWithLazyDBRefsAndConstructorCreation(null, one,
+				null, null);
+
+		template.save(source);
+
+		DocumentWithLazyDBRefsAndConstructorCreation target = template.findOne(query(where("id").is(source.id)),
+				DocumentWithLazyDBRefsAndConstructorCreation.class);
+
+		assertThat(target.lazyDbRefProperty, instanceOf(LazyLoadingProxy.class));
+		assertThat(target.lazyDbRefProperty, is(one));
+	}
+
+	@Test // DATAMONGO-2004
+	public void shouldFetchMapOfLazyReferencesWithConstructorCreationCorrectly() {
+
+		Sample one = new Sample("1", "jon snow");
+		Sample two = new Sample("2", "tyrion lannister");
+
+		template.save(one);
+		template.save(two);
+
+		Map<String, Sample> map = new LinkedHashMap<>();
+		map.put("tyrion", two);
+		map.put("jon", one);
+
+		DocumentWithLazyDBRefsAndConstructorCreation source = new DocumentWithLazyDBRefsAndConstructorCreation(null, null,
+				null, map);
+
+		template.save(source);
+
+		DocumentWithLazyDBRefsAndConstructorCreation target = template.findOne(query(where("id").is(source.id)),
+				DocumentWithLazyDBRefsAndConstructorCreation.class);
+
+		assertThat(target.lazyDbRefAnnotatedMap, instanceOf(LazyLoadingProxy.class));
+		assertThat(target.lazyDbRefAnnotatedMap.values(), contains(two, one));
+	}
+
+	@Test // DATAMONGO-2004
+	public void shouldFetchListOfLazyReferencesWithConstructorCreationCorrectly() {
+
+		Sample one = new Sample("1", "jon snow");
+		Sample two = new Sample("2", "tyrion lannister");
+
+		template.save(one);
+		template.save(two);
+
+		List<Sample> list = Arrays.asList(two, one);
+
+		DocumentWithLazyDBRefsAndConstructorCreation source = new DocumentWithLazyDBRefsAndConstructorCreation(null, null,
+				list, null);
+
+		template.save(source);
+
+		DocumentWithLazyDBRefsAndConstructorCreation target = template.findOne(query(where("id").is(source.id)),
+				DocumentWithLazyDBRefsAndConstructorCreation.class);
+
+		assertThat(target.lazyDbRefAnnotatedList, instanceOf(LazyLoadingProxy.class));
+		assertThat(target.lazyDbRefAnnotatedList, contains(two, one));
+	}
+
 	@Test // DATAMONGO-1513
 	@DirtiesContext
 	public void populatesIdsAddedByEventListener() {
@@ -3590,6 +3657,29 @@ public class MongoTemplateTests {
 		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) public Map<String, Sample> lazyDbRefAnnotatedMap;
 	}
 
+	@Data
+	static class DocumentWithLazyDBRefsAndConstructorCreation {
+
+		@Id public final String id;
+
+		public DocumentWithLazyDBRefsAndConstructorCreation(String id, Sample lazyDbRefProperty,
+				List<Sample> lazyDbRefAnnotatedList, Map<String, Sample> lazyDbRefAnnotatedMap) {
+			this.id = id;
+			this.lazyDbRefProperty = lazyDbRefProperty;
+			this.lazyDbRefAnnotatedList = lazyDbRefAnnotatedList;
+			this.lazyDbRefAnnotatedMap = lazyDbRefAnnotatedMap;
+		}
+
+		@org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) //
+		public final Sample lazyDbRefProperty;
+
+		@Field("lazy_db_ref_list") @org.springframework.data.mongodb.core.mapping.DBRef(lazy = true) //
+		public final List<Sample> lazyDbRefAnnotatedList;
+
+		@Field("lazy_db_ref_map") @org.springframework.data.mongodb.core.mapping.DBRef(
+				lazy = true) public final Map<String, Sample> lazyDbRefAnnotatedMap;
+	}
+
 	@EqualsAndHashCode
 	static class DocumentWithCollection {