Enhance support for linking entities.

Add initial support for an alternative to the existing DBRef scenario.
The enhancement allows to store and retrieve linked entites via their id or a customizable lookup query.

Author: christophstrobl

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

@@ -35,7 +35,7 @@
  * @author Mark Paluch
  * @since 1.4
  */
-public interface DbRefResolver {
+public interface DbRefResolver extends ReferenceResolver {
 
 	/**
 	 * Resolves the given {@link DBRef} into an object of the given {@link MongoPersistentProperty}'s type. The method

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

@@ -46,6 +46,7 @@
 import org.springframework.data.mongodb.LazyLoadingException;
 import org.springframework.data.mongodb.MongoDatabaseFactory;
 import org.springframework.data.mongodb.MongoDatabaseUtils;
+import org.springframework.data.mongodb.core.convert.ReferenceLoader.ReferenceFilter;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.lang.Nullable;
 import org.springframework.objenesis.ObjenesisStd;
@@ -67,7 +68,7 @@
  * @author Mark Paluch
  * @since 1.4
  */
-public class DefaultDbRefResolver implements DbRefResolver {
+public class DefaultDbRefResolver extends DefaultReferenceResolver implements DbRefResolver, ReferenceResolver {
 
 	private static final Logger LOGGER = LoggerFactory.getLogger(DefaultDbRefResolver.class);
 
@@ -82,6 +83,8 @@ public class DefaultDbRefResolver implements DbRefResolver {
 	 */
 	public DefaultDbRefResolver(MongoDatabaseFactory mongoDbFactory) {
 
+		super(new DefaultReferenceLoader(mongoDbFactory));
+
 		Assert.notNull(mongoDbFactory, "MongoDbFactory translator must not be null!");
 
 		this.mongoDbFactory = mongoDbFactory;
@@ -114,17 +117,7 @@ public Object resolveDbRef(MongoPersistentProperty property, @Nullable DBRef dbr
 	 */
 	@Override
 	public Document fetch(DBRef dbRef) {
-
-		MongoCollection<Document> mongoCollection = getCollection(dbRef);
-
-		if (LOGGER.isTraceEnabled()) {
-			LOGGER.trace("Fetching DBRef '{}' from {}.{}.", dbRef.getId(),
-					StringUtils.hasText(dbRef.getDatabaseName()) ? dbRef.getDatabaseName()
-							: mongoCollection.getNamespace().getDatabaseName(),
-					dbRef.getCollectionName());
-		}
-
-		return mongoCollection.find(Filters.eq("_id", dbRef.getId())).first();
+		return getReferenceLoader().fetch(ReferenceFilter.singleReferenceFilter(Filters.eq("_id", dbRef.getId())), ReferenceContext.fromDBRef(dbRef));
 	}
 
 	/*
@@ -164,9 +157,9 @@ public List<Document> bulkFetch(List<DBRef> refs) {
 					databaseSource.getCollectionName());
 		}
 
-		List<Document> result = mongoCollection //
-				.find(new Document("_id", new Document("$in", ids))) //
-				.into(new ArrayList<>());
+		List<Document> result = getReferenceLoader()
+				.bulkFetch(ReferenceFilter.referenceFilter(new Document("_id", new Document("$in", ids))), ReferenceContext.fromDBRef(refs.iterator().next()))
+				.collect(Collectors.toList());
 
 		return ids.stream() //
 				.flatMap(id -> documentWithId(id, result)) //
@@ -504,4 +497,10 @@ protected MongoCollection<Document> getCollection(DBRef dbref) {
 		return MongoDatabaseUtils.getDatabase(dbref.getDatabaseName(), mongoDbFactory)
 				.getCollection(dbref.getCollectionName(), Document.class);
 	}
+
+	protected MongoCollection<Document> getCollection(ReferenceContext context) {
+
+		return MongoDatabaseUtils.getDatabase(context.database, mongoDbFactory).getCollection(context.collection,
+				Document.class);
+	}
 }

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultReferenceLoader.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2021 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
+ *
+ *      https://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 java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
+import org.bson.Document;
+import org.bson.conversions.Bson;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.data.mongodb.MongoDatabaseFactory;
+import org.springframework.data.mongodb.MongoDatabaseUtils;
+import org.springframework.data.mongodb.core.convert.ReferenceResolver.ReferenceContext;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
+
+import com.mongodb.client.FindIterable;
+import com.mongodb.client.MongoCollection;
+
+/**
+ * @author Christoph Strobl
+ */
+public class DefaultReferenceLoader implements ReferenceLoader {
+
+	private static final Logger LOGGER = LoggerFactory.getLogger(DefaultReferenceLoader.class);
+
+	private final MongoDatabaseFactory mongoDbFactory;
+
+	public DefaultReferenceLoader(MongoDatabaseFactory mongoDbFactory) {
+
+		Assert.notNull(mongoDbFactory, "MongoDbFactory translator must not be null!");
+
+		this.mongoDbFactory = mongoDbFactory;
+	}
+
+	@Override
+	public Stream<Document> bulkFetch(ReferenceFilter filter, ReferenceContext context) {
+
+		MongoCollection<Document> collection = getCollection(context);
+
+		if (LOGGER.isTraceEnabled()) {
+			LOGGER.trace("Bulk fetching {} from {}.{}.", filter,
+					StringUtils.hasText(context.getDatabase()) ? context.getDatabase()
+							: collection.getNamespace().getDatabaseName(),
+					context.getCollection());
+		}
+
+		return filter.apply(collection);
+	}
+
+	protected MongoCollection<Document> getCollection(ReferenceContext context) {
+
+		return MongoDatabaseUtils.getDatabase(context.database, mongoDbFactory).getCollection(context.collection,
+				Document.class);
+	}
+}

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultReferenceResolver.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2021 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
+ *
+ *      https://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 java.util.function.BiFunction;
+import java.util.stream.Stream;
+
+import org.bson.Document;
+import org.bson.conversions.Bson;
+import org.springframework.data.mongodb.core.convert.ReferenceLoader.ReferenceFilter;
+import org.springframework.data.mongodb.core.mapping.DocumentReference;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.lang.Nullable;
+
+/**
+ * @author Christoph Strobl
+ */
+public class DefaultReferenceResolver implements ReferenceResolver {
+
+	private final ReferenceLoader referenceLoader;
+
+	public DefaultReferenceResolver(ReferenceLoader referenceLoader) {
+		this.referenceLoader = referenceLoader;
+	}
+
+	@Override
+	public ReferenceLoader getReferenceLoader() {
+		return referenceLoader;
+	}
+
+	@Nullable
+	@Override
+	public Object resolveReference(MongoPersistentProperty property, Object source, ReferenceReader referenceReader,
+			BiFunction<ReferenceContext, ReferenceFilter, Stream<Document>> lookupFunction) {
+
+		if (isLazyReference(property)) {
+			return createLazyLoadingProxy(property, source, referenceReader, lookupFunction);
+		}
+
+		return referenceReader.readReference(property, source, lookupFunction);
+	}
+
+	private Object createLazyLoadingProxy(MongoPersistentProperty property, Object source,
+			ReferenceReader referenceReader, BiFunction<ReferenceContext, ReferenceFilter, Stream<Document>> lookupFunction) {
+		return new LazyLoadingProxyGenerator(referenceReader).createLazyLoadingProxy(property, source, lookupFunction);
+	}
+
+	protected boolean isLazyReference(MongoPersistentProperty property) {
+
+		if (property.findAnnotation(DocumentReference.class) != null) {
+			return property.findAnnotation(DocumentReference.class).lazy();
+		}
+
+		return property.getDBRef() != null && property.getDBRef().lazy();
+	}
+}