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.

Author: mp911de

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

@@ -58,6 +58,14 @@ public DocumentAccessor(Bson document) {
 		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 @@ private <S extends Object> S read(TypeInformation<S> type, @Nullable Bson bson,
 	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 @@ private Object removeTypeInfo(Object object, boolean recursively) {
 	 */
 	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 MongoDbPropertyValueProvider(DocumentAccessor accessor, SpELExpressionEva
 		 * @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 <T> T getPropertyValue(MongoPersistentProperty property) {
 		}
 	}
 
+	/**
+	 * {@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

@@ -2935,8 +2935,8 @@ public void resolvesCyclicDBRefCorrectly() {
 		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";
@@ -2949,7 +2949,7 @@ public void shouldReuseAlreadyResolvedLazyLoadedDBRefWhenUsedAsPersistenceConstr
 
 		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());
 	}
 
@@ -2972,8 +2972,8 @@ public void shouldNotReuseLazyLoadedDBRefWhenTypeUsedInPersistenceConstrcutorBut
 		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";
@@ -2991,7 +2991,7 @@ public void shouldRespectParamterValueWhenAttemptingToReuseLazyLoadedDBRefUsedIn
 
 		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));
 	}
 
@@ -3250,6 +3250,73 @@ public void shouldFetchMapOfLazyReferencesCorrectly() {
 		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() {
@@ -3457,6 +3524,29 @@ static class DocumentWithDBRefCollection {
 		@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 {