Map collection and fields for $lookup aggregation against type.

christophstrobl · christophstrobl · commit 3a9b3694a3c6 · 2023-06-12T11:35:19.000+02:00
This commit enables using a type parameter to define the from collection of a lookup aggregation stage. In doing so we can derive the target collection name from the type and use the given information to also map the from field against the domain object to so that the user is able to operate on property names instead of the target db field name.
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperationContext.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/AggregationOperationContext.java
@@ -23,6 +23,7 @@
 import org.bson.codecs.configuration.CodecRegistry;
 import org.springframework.beans.BeanUtils;
 import org.springframework.data.mongodb.CodecRegistryProvider;
+import org.springframework.data.mongodb.MongoCollectionUtils;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -78,6 +79,29 @@ default Document getMappedObject(Document document) {
 	 */
 	FieldReference getReference(String name);
 
+	/**
+	 * Obtain the target field name for a given field/type combination.
+	 * 
+	 * @param type The type containing the field.
+	 * @param field The property/field name
+	 * @return never {@literal null}.
+	 * @since 4.2
+	 */
+	default String getMappedFieldName(Class<?> type, String field) {
+		return field;
+	}
+
+	/**
+	 * Obtain the collection name for a given {@link Class type} combination.
+	 *
+	 * @param type
+	 * @return never {@literal null}.
+	 * @since 4.2
+	 */
+	default String getCollection(Class<?> type) {
+		return MongoCollectionUtils.getPreferredCollectionName(type);
+	}
+
 	/**
 	 * Returns the {@link Fields} exposed by the type. May be a {@literal class} or an {@literal interface}. The default
 	 * implementation uses {@link BeanUtils#getPropertyDescriptors(Class) property descriptors} discover fields from a
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java
@@ -39,7 +39,7 @@
  */
 public class LookupOperation implements FieldsExposingAggregationOperation, InheritsFieldsAggregationOperation {
 
-	private final String from;
+	private Object from;
 
 	@Nullable //
 	private final Field localField;
@@ -97,6 +97,22 @@ public LookupOperation(String from, @Nullable Let let, AggregationPipeline pipel
 	 */
 	public LookupOperation(String from, @Nullable Field localField, @Nullable Field foreignField, @Nullable Let let,
 			@Nullable AggregationPipeline pipeline, Field as) {
+		this((Object) from, localField, foreignField, let, pipeline, as);
+	}
+
+	/**
+	 * Creates a new {@link LookupOperation} for the given combination of {@link Field}s and {@link AggregationPipeline
+	 * pipeline}.
+	 *
+	 * @param from must not be {@literal null}. Can be eiter the target collection name or a {@link Class}.
+	 * @param localField can be {@literal null} if {@literal pipeline} is present.
+	 * @param foreignField can be {@literal null} if {@literal pipeline} is present.
+	 * @param let can be {@literal null} if {@literal localField} and {@literal foreignField} are present.
+	 * @param as must not be {@literal null}.
+	 * @since 4.1
+	 */
+	private LookupOperation(Object from, @Nullable Field localField, @Nullable Field foreignField, @Nullable Let let,
+			@Nullable AggregationPipeline pipeline, Field as) {
 
 		Assert.notNull(from, "From must not be null");
 		if (pipeline == null) {
@@ -125,12 +141,14 @@ public Document toDocument(AggregationOperationContext context) {
 
 		Document lookupObject = new Document();
 
-		lookupObject.append("from", from);
+		lookupObject.append("from", getCollectionName(context));
+		
 		if (localField != null) {
 			lookupObject.append("localField", localField.getTarget());
 		}
+		
 		if (foreignField != null) {
-			lookupObject.append("foreignField", foreignField.getTarget());
+			lookupObject.append("foreignField", getForeignFieldName(context));
 		}
 		if (let != null) {
 			lookupObject.append("let", let.toDocument(context).get("$let", Document.class).get("vars"));
@@ -144,6 +162,16 @@ public Document toDocument(AggregationOperationContext context) {
 		return new Document(getOperator(), lookupObject);
 	}
 
+	String getCollectionName(AggregationOperationContext context) {
+		return from instanceof Class<?> type ? context.getCollection(type) : from.toString();
+	}
+
+	String getForeignFieldName(AggregationOperationContext context) {
+
+		return from instanceof Class<?> type ? context.getMappedFieldName(type, foreignField.getTarget())
+				: foreignField.getTarget();
+	}
+
 	@Override
 	public String getOperator() {
 		return "$lookup";
@@ -158,16 +186,28 @@ public static FromBuilder newLookup() {
 		return new LookupOperationBuilder();
 	}
 
-	public static interface FromBuilder {
+	public interface FromBuilder {
 
 		/**
 		 * @param name the collection in the same database to perform the join with, must not be {@literal null} or empty.
 		 * @return never {@literal null}.
 		 */
 		LocalFieldBuilder from(String name);
+
+		/**
+		 * Use the given type to determine name of the foreign collection and map
+		 * {@link ForeignFieldBuilder#foreignField(String)} against it to consider eventually present
+		 * {@link org.springframework.data.mongodb.core.mapping.Field} annotations.
+		 * 
+		 * @param type the type of the target collection in the same database to perform the join with, must not be
+		 *          {@literal null}.
+		 * @return never {@literal null}.
+		 * @since 4.2
+		 */
+		LocalFieldBuilder from(Class<?> type);
 	}
 
-	public static interface LocalFieldBuilder extends PipelineBuilder {
+	public interface LocalFieldBuilder extends PipelineBuilder {
 
 		/**
 		 * @param name the field from the documents input to the {@code $lookup} stage, must not be {@literal null} or
@@ -177,7 +217,7 @@ public static interface LocalFieldBuilder extends PipelineBuilder {
 		ForeignFieldBuilder localField(String name);
 	}
 
-	public static interface ForeignFieldBuilder {
+	public interface ForeignFieldBuilder {
 
 		/**
 		 * @param name the field from the documents in the {@code from} collection, must not be {@literal null} or empty.
@@ -246,7 +286,7 @@ default AsBuilder pipeline(AggregationOperation... stages) {
 		LookupOperation as(String name);
 	}
 
-	public static interface AsBuilder extends PipelineBuilder {
+	public interface AsBuilder extends PipelineBuilder {
 
 		/**
 		 * @param name the name of the new array field to add to the input documents, must not be {@literal null} or empty.
@@ -264,7 +304,7 @@ public static interface AsBuilder extends PipelineBuilder {
 	public static final class LookupOperationBuilder
 			implements FromBuilder, LocalFieldBuilder, ForeignFieldBuilder, AsBuilder {
 
-		private @Nullable String from;
+		private @Nullable Object from;
 		private @Nullable Field localField;
 		private @Nullable Field foreignField;
 		private @Nullable ExposedField as;
@@ -288,6 +328,14 @@ public LocalFieldBuilder from(String name) {
 			return this;
 		}
 
+		@Override
+		public LocalFieldBuilder from(Class<?> type) {
+
+			Assert.notNull(type, "'From' must not be null");
+			from = type;
+			return this;
+		}
+
 		@Override
 		public AsBuilder foreignField(String name) {
 
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContext.java
@@ -92,6 +92,20 @@ public FieldReference getReference(String name) {
 		return getReferenceFor(field(name));
 	}
 
+	@Override
+	public String getCollection(Class<?> type) {
+
+		MongoPersistentEntity<?> persistentEntity = mappingContext.getPersistentEntity(type);
+		return persistentEntity != null ? persistentEntity.getCollection() : AggregationOperationContext.super.getCollection(type);
+	}
+
+	@Override
+	public String getMappedFieldName(Class<?> type, String field) {
+
+		PersistentPropertyPath<MongoPersistentProperty> persistentPropertyPath = mappingContext.getPersistentPropertyPath(field, type);
+		return persistentPropertyPath.getLeafProperty().getFieldName();
+	}
+
 	@Override
 	public Fields getFields(Class<?> type) {
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTestUtils.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTestUtils.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2023. 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
+ *
+ *      http://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.
+ */
+
+/*
+ * Copyright 2023 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
+ *
+ *      http://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.aggregation;
+
+import org.springframework.data.mapping.context.MappingContext;
+import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
+import org.springframework.data.mongodb.core.convert.NoOpDbRefResolver;
+import org.springframework.data.mongodb.core.convert.QueryMapper;
+import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.data.mongodb.test.util.MongoTestMappingContext;
+
+/**
+ * @author Christoph Strobl
+ * @since 2023/06
+ */
+public final class AggregationTestUtils {
+
+	public static AggregationContextBuilder<TypeBasedAggregationOperationContext> strict(Class<?> type) {
+
+		AggregationContextBuilder<AggregationOperationContext> builder = new AggregationContextBuilder<>();
+		builder.strict = true;
+		return builder.forType(type);
+	}
+
+	public static AggregationContextBuilder<TypeBasedAggregationOperationContext> relaxed(Class<?> type) {
+
+		AggregationContextBuilder<AggregationOperationContext> builder = new AggregationContextBuilder<>();
+		builder.strict = false;
+		return builder.forType(type);
+	}
+
+	public static class AggregationContextBuilder<T extends AggregationOperationContext> {
+
+		Class<?> targetType;
+		MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
+		QueryMapper queryMapper;
+		boolean strict;
+
+		public AggregationContextBuilder<TypeBasedAggregationOperationContext> forType(Class<?> type) {
+
+			this.targetType = type;
+			return (AggregationContextBuilder<TypeBasedAggregationOperationContext>) this;
+		}
+
+		public AggregationContextBuilder<T> using(
+				MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext) {
+
+			this.mappingContext = mappingContext;
+			return this;
+		}
+
+		public AggregationContextBuilder<T> using(QueryMapper queryMapper) {
+
+			this.queryMapper = queryMapper;
+			return this;
+		}
+
+		public T ctx() {
+			//
+			if (targetType == null) {
+				return (T) Aggregation.DEFAULT_CONTEXT;
+			}
+
+			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty>  ctx = mappingContext != null ? mappingContext : MongoTestMappingContext.newTestContext().init();
+			QueryMapper qm = queryMapper != null ? queryMapper
+					: new QueryMapper(new MappingMongoConverter(NoOpDbRefResolver.INSTANCE, ctx));
+			return (T) (strict ? new TypeBasedAggregationOperationContext(targetType, ctx, qm)
+					: new RelaxedTypeBasedAggregationOperationContext(targetType, ctx, qm));
+		}
+	}
+}
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/LookupOperationUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/LookupOperationUnitTests.java
@@ -25,6 +25,7 @@
 import org.bson.Document;
 import org.junit.jupiter.api.Test;
 import org.springframework.data.mongodb.core.DocumentTestUtils;
+import org.springframework.data.mongodb.core.mapping.Field;
 import org.springframework.data.mongodb.core.query.Criteria;
 
 /**
@@ -92,7 +93,7 @@ private Document extractDocumentFromLookupOperation(LookupOperation lookupOperat
 
 	@Test // DATAMONGO-1326
 	public void builderRejectsNullFromField() {
-		assertThatIllegalArgumentException().isThrownBy(() -> LookupOperation.newLookup().from(null));
+		assertThatIllegalArgumentException().isThrownBy(() -> LookupOperation.newLookup().from((String) null));
 	}
 
 	@Test // DATAMONGO-1326
@@ -195,10 +196,10 @@ void buildsLookupWithJustPipeline() {
 	void buildsLookupWithLocalAndForeignFieldAsWellAsLetAndPipeline() {
 
 		LookupOperation lookupOperation = Aggregation.lookup().from("restaurants") //
-				.localField("restaurant_name")
-				.foreignField("name")
+				.localField("restaurant_name") //
+				.foreignField("name") //
 				.let(newVariable("orders_drink").forField("drink")) //
-				.pipeline(match(ctx -> new Document("$expr", new Document("$in", List.of("$$orders_drink", "$beverages")))))
+				.pipeline(match(ctx -> new Document("$expr", new Document("$in", List.of("$$orders_drink", "$beverages"))))) //
 				.as("matches");
 
 		assertThat(lookupOperation.toDocument(Aggregation.DEFAULT_CONTEXT)).isEqualTo("""
@@ -216,4 +217,54 @@ void buildsLookupWithLocalAndForeignFieldAsWellAsLetAndPipeline() {
 				}}
 				""");
 	}
+
+	@Test // GH-4379
+	void unmappedLookupWithFromExtractedFromType() {
+
+		LookupOperation lookupOperation = Aggregation.lookup().from(Restaurant.class) //
+				.localField("restaurant_name") //
+				.foreignField("name") //
+				.as("restaurants");
+
+		assertThat(lookupOperation.toDocument(Aggregation.DEFAULT_CONTEXT)).isEqualTo("""
+				{ $lookup:
+					{
+						from: "restaurant",
+						localField: "restaurant_name",
+						foreignField: "name",
+						as: "restaurants"
+					}
+				}}
+				""");
+	}
+
+	@Test // GH-4379
+	void mappedLookupWithFromExtractedFromType() {
+
+		LookupOperation lookupOperation = Aggregation.lookup().from(Restaurant.class) //
+				.localField("restaurant_name") //
+				.foreignField("name") //
+				.as("restaurants");
+
+
+		assertThat(lookupOperation.toDocument(AggregationTestUtils.strict(Restaurant.class).ctx())).isEqualTo("""
+				{ $lookup:
+					{
+						from: "sites",
+						localField: "restaurant_name",
+						foreignField: "rs_name",
+						as: "restaurants"
+					}
+				}}
+				""");
+	}
+
+	@org.springframework.data.mongodb.core.mapping.Document("sites")
+	static class Restaurant {
+
+		String id;
+
+		@Field("rs_name") //
+		String name;
+	}
 }
