DATAMONGO-1326 - Support field inheritance for $lookup aggregation operator.

We now distinguish between aggregation operations that replace fields in the aggregation pipeline and those which inherit fields from previous operations. InheritsFieldsAggregationOperation is a nested interface of FieldsExposingAggregationOperation is a marker to lookup fields along the aggregation context chain. Added unit and integration tests. Mention lookup operator in docs.

Original pull request: #344.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/Aggregation.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.
@@ -26,6 +26,7 @@ import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
 import org.springframework.data.mongodb.core.aggregation.Fields.AggregationField;
+import org.springframework.data.mongodb.core.aggregation.FieldsExposingAggregationOperation.InheritsFieldsAggregationOperation;
 import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.data.mongodb.core.query.SerializationUtils;
@@ -41,6 +42,7 @@ import com.mongodb.DBObject;
  * @author Tobias Trelle
  * @author Thomas Darimont
  * @author Oliver Gierke
+ * @author Mark Paluch
  * @author Alessio Fachechi
  * @since 1.3
  */
@@ -362,7 +364,12 @@ public class Aggregation {
 			if (operation instanceof FieldsExposingAggregationOperation) {
 
 				FieldsExposingAggregationOperation exposedFieldsOperation = (FieldsExposingAggregationOperation) operation;
-				context = new ExposedFieldsAggregationOperationContext(exposedFieldsOperation.getFields(), rootContext);
+
+				if (operation instanceof InheritsFieldsAggregationOperation) {
+					context = new InheritingExposedFieldsAggregationOperationContext(exposedFieldsOperation.getFields(), context);
+				} else {
+					context = new ExposedFieldsAggregationOperationContext(exposedFieldsOperation.getFields(), context);
+				}
 			}
 		}
 

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/ExposedFieldsAggregationOperationContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2014 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.
@@ -27,6 +27,7 @@ import com.mongodb.DBObject;
  *
  * @author Thomas Darimont
  * @author Oliver Gierke
+ * @author Mark Paluch
  * @since 1.4
  */
 class ExposedFieldsAggregationOperationContext implements AggregationOperationContext {
@@ -89,6 +90,22 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 
 		Assert.notNull(name, "Name must not be null!");
 
+		FieldReference exposedField = resolveExposedField(field, name);
+		if (exposedField != null) {
+			return exposedField;
+		}
+
+		throw new IllegalArgumentException(String.format("Invalid reference '%s'!", name));
+	}
+
+	/**
+	 * Resolves a {@link field}/{@link name} for a {@link FieldReference} if possible.
+	 *
+	 * @param field may be {@literal null}
+	 * @param name must not be {@literal null}
+	 * @return the resolved reference or {@literal null}
+	 */
+	protected FieldReference resolveExposedField(Field field, String name) {
 		ExposedField exposedField = exposedFields.getField(name);
 
 		if (exposedField != null) {
@@ -112,7 +129,6 @@ class ExposedFieldsAggregationOperationContext implements AggregationOperationCo
 				return new FieldReference(new ExposedField(name, true));
 			}
 		}
-
-		throw new IllegalArgumentException(String.format("Invalid reference '%s'!", name));
+		return null;
 	}
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/FieldsExposingAggregationOperation.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 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.
@@ -16,17 +16,28 @@
 package org.springframework.data.mongodb.core.aggregation;
 
 /**
- * {@link AggregationOperation} that exposes new {@link ExposedFields} that can be used for later aggregation pipeline
- * {@code AggregationOperation}s.
- * 
+ * {@link AggregationOperation} that exposes {@link ExposedFields} that can be used for later aggregation pipeline
+ * {@code AggregationOperation}s. A {@link FieldsExposingAggregationOperation} implementing the
+ * {@link InheritsFieldsAggregationOperation} will expose fields from its parent operations. Not implementing
+ * {@link InheritsFieldsAggregationOperation} will replace existing exposed fields.
+ *
  * @author Thomas Darimont
+ * @author Mark Paluch
  */
 public interface FieldsExposingAggregationOperation extends AggregationOperation {
 
 	/**
 	 * Returns the fields exposed by the {@link AggregationOperation}.
-	 * 
+	 *
 	 * @return will never be {@literal null}.
 	 */
 	ExposedFields getFields();
+
+	/**
+	 * Marker interface for {@link AggregationOperation} that inherits fields from previous operations.
+	 */
+	static interface InheritsFieldsAggregationOperation extends FieldsExposingAggregationOperation {
+
+	}
+
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/InheritingExposedFieldsAggregationOperationContext.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 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.
+ * 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.mongodb.core.aggregation.ExposedFields.FieldReference;
+import org.springframework.util.Assert;
+
+/**
+ * {@link ExposedFieldsAggregationOperationContext} that inherits fields from its parent
+ * {@link AggregationOperationContext}.
+ *
+ * @author Mark Paluch
+ */
+class InheritingExposedFieldsAggregationOperationContext extends ExposedFieldsAggregationOperationContext {
+
+	private final AggregationOperationContext previousContext;
+
+	/**
+	 * Creates a new {@link ExposedFieldsAggregationOperationContext} from the given {@link ExposedFields}. Uses the given
+	 * {@link AggregationOperationContext} to perform a mapping to mongo types if necessary.
+	 *
+	 * @param exposedFields must not be {@literal null}.
+	 * @param previousContext must not be {@literal null}.
+	 */
+	public InheritingExposedFieldsAggregationOperationContext(ExposedFields exposedFields,
+			AggregationOperationContext previousContext) {
+
+		super(exposedFields, previousContext);
+		Assert.notNull(previousContext, "PreviousContext must not be null!");
+		this.previousContext = previousContext;
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.aggregation.ExposedFieldsAggregationOperationContext#resolveExposedField(org.springframework.data.mongodb.core.aggregation.Field, java.lang.String)
+	 */
+	@Override
+	protected FieldReference resolveExposedField(Field field, String name) {
+
+		FieldReference fieldReference = super.resolveExposedField(field, name);
+		if (fieldReference != null) {
+			return fieldReference;
+		}
+
+		if (field != null) {
+			return previousContext.getReference(field);
+		}
+
+		return previousContext.getReference(name);
+	}
+}

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/aggregation/LookupOperation.java

@@ -16,6 +16,7 @@
 package org.springframework.data.mongodb.core.aggregation;
 
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
+import org.springframework.data.mongodb.core.aggregation.FieldsExposingAggregationOperation.InheritsFieldsAggregationOperation;
 import org.springframework.util.Assert;
 
 import com.mongodb.BasicDBObject;
@@ -27,10 +28,11 @@ import com.mongodb.DBObject;
  *
  * @author Alessio Fachechi
  * @author Christoph Strobl
+ * @author Mark Paluch
  * @see http://docs.mongodb.org/manual/reference/aggregation/lookup/#stage._S_lookup
  * @since 1.9
  */
-public class LookupOperation implements FieldsExposingAggregationOperation {
+public class LookupOperation implements FieldsExposingAggregationOperation, InheritsFieldsAggregationOperation {
 
 	private Field from;
 	private Field localField;
@@ -100,7 +102,7 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 	public static interface FromBuilder {
 
 		/**
-		 * @param name
+		 * @param name the collection in the same database to perform the join with, must not be {@literal null} or empty.
 		 * @return
 		 */
 		LocalFieldBuilder from(String name);
@@ -109,7 +111,8 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 	public static interface LocalFieldBuilder {
 
 		/**
-		 * @param name
+		 * @param name the field from the documents input to the {@code $lookup} stage, must not be {@literal null} or
+		 *          empty.
 		 * @return
 		 */
 		ForeignFieldBuilder localField(String name);
@@ -118,7 +121,7 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 	public static interface ForeignFieldBuilder {
 
 		/**
-		 * @param name
+		 * @param name the field from the documents in the {@code from} collection, must not be {@literal null} or empty.
 		 * @return
 		 */
 		AsBuilder foreignField(String name);
@@ -127,7 +130,7 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 	public static interface AsBuilder {
 
 		/**
-		 * @param name
+		 * @param name the name of the new array field to add to the input documents, must not be {@literal null} or empty.
 		 * @return
 		 */
 		LookupOperation as(String name);
@@ -135,14 +138,14 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 
 	/**
 	 * Builder for fluent {@link LookupOperation} creation.
-	 * 
+	 *
 	 * @author Christoph Strobl
 	 * @since 1.9
 	 */
 	public static final class LookupOperationBuilder
 			implements FromBuilder, LocalFieldBuilder, ForeignFieldBuilder, AsBuilder {
 
-		private LookupOperation lookupOperation;
+		private final LookupOperation lookupOperation;
 
 		private LookupOperationBuilder() {
 			this.lookupOperation = new LookupOperation();
@@ -150,7 +153,7 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 
 		/**
 		 * Creates new builder for {@link LookupOperation}.
-		 * 
+		 *
 		 * @return never {@literal null}.
 		 */
 		public static FromBuilder newBuilder() {
@@ -170,7 +173,8 @@ public class LookupOperation implements FieldsExposingAggregationOperation {
 
 			Assert.hasText(name, "'As' must not be null or empty!");
 			lookupOperation.as = new ExposedField(Fields.field(name), true);
-			return null;
+			return new LookupOperation(lookupOperation.from, lookupOperation.localField, lookupOperation.foreignField,
+					lookupOperation.as);
 		}
 
 		@Override

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/AggregationTests.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.
@@ -21,6 +21,7 @@ import static org.junit.Assume.*;
 import static org.springframework.data.domain.Sort.Direction.*;
 import static org.springframework.data.mongodb.core.aggregation.Aggregation.*;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
+import static org.springframework.data.mongodb.test.util.IsBsonObject.*;
 
 import java.io.BufferedInputStream;
 import java.text.ParseException;
@@ -76,6 +77,7 @@ import com.mongodb.util.JSON;
  * @author Thomas Darimont
  * @author Oliver Gierke
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration("classpath:infrastructure.xml")
@@ -85,6 +87,7 @@ public class AggregationTests {
 	private static final Logger LOGGER = LoggerFactory.getLogger(AggregationTests.class);
 	private static final Version TWO_DOT_FOUR = new Version(2, 4);
 	private static final Version TWO_DOT_SIX = new Version(2, 6);
+	private static final Version THREE_DOT_TWO = new Version(3, 2);
 
 	private static boolean initialized = false;
 
@@ -1068,6 +1071,76 @@ public class AggregationTests {
 		assertThat(result.get("totalValue"), is(equalTo((Object) 100.0)));
 	}
 
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void shouldLookupPeopleCorectly() {
+
+		assumeTrue(mongoVersion.isGreaterThanOrEqualTo(THREE_DOT_TWO));
+
+		createUsersWithReferencedPersons();
+
+		TypedAggregation<User> agg = newAggregation(User.class, //
+				lookup("person", "_id", "firstname", "linkedPerson"), //
+				sort(ASC, "id"));
+
+		AggregationResults<DBObject> results = mongoTemplate.aggregate(agg, User.class, DBObject.class);
+
+		List<DBObject> mappedResults = results.getMappedResults();
+
+		DBObject firstItem = mappedResults.get(0);
+
+		assertThat(firstItem, isBsonObject().containing("_id", "u1"));
+		assertThat(firstItem, isBsonObject().containing("linkedPerson.[0].firstname", "u1"));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void shouldGroupByAndLookupPeopleCorectly() {
+
+		assumeTrue(mongoVersion.isGreaterThanOrEqualTo(THREE_DOT_TWO));
+
+		createUsersWithReferencedPersons();
+
+		TypedAggregation<User> agg = newAggregation(User.class, //
+				group().min("id").as("foreignKey"), //
+				lookup("person", "foreignKey", "firstname", "linkedPerson"), //
+				sort(ASC, "foreignKey", "linkedPerson.firstname"));
+
+		AggregationResults<DBObject> results = mongoTemplate.aggregate(agg, User.class, DBObject.class);
+
+		List<DBObject> mappedResults = results.getMappedResults();
+
+		DBObject firstItem = mappedResults.get(0);
+
+		assertThat(firstItem, isBsonObject().containing("foreignKey", "u1"));
+		assertThat(firstItem, isBsonObject().containing("linkedPerson.[0].firstname", "u1"));
+	}
+
+	private void createUsersWithReferencedPersons() {
+
+		mongoTemplate.dropCollection(User.class);
+		mongoTemplate.dropCollection(Person.class);
+
+		User user1 = new User("u1");
+		User user2 = new User("u2");
+		User user3 = new User("u3");
+
+		mongoTemplate.save(user1);
+		mongoTemplate.save(user2);
+		mongoTemplate.save(user3);
+
+		Person person1 = new Person("u1", "User 1");
+		Person person2 = new Person("u2", "User 2");
+
+		mongoTemplate.save(person1);
+		mongoTemplate.save(person2);
+		mongoTemplate.save(user3);
+	}
+
 	private void assertLikeStats(LikeStats like, String id, long count) {
 
 		assertThat(like, is(notNullValue()));

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

@@ -29,6 +29,7 @@ import com.mongodb.DBObject;
  *
  * @author Alessio Fachechi
  * @author Christoph Strobl
+ * @author Mark Paluch
  */
 public class LookupOperationUnitTests {
 
@@ -100,4 +101,66 @@ public class LookupOperationUnitTests {
 		DBObject lookupClause = DBObjectTestUtils.getAsDBObject(dbObject, "$lookup");
 		return lookupClause;
 	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void builderRejectsNullFromField() {
+		LookupOperation.newLookup().from(null);
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void builderRejectsNullLocalField() {
+		LookupOperation.newLookup().from("a").localField(null);
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void builderRejectsNullForeignField() {
+		LookupOperation.newLookup().from("a").localField("b").foreignField(null);
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void builderRejectsNullAsField() {
+		LookupOperation.newLookup().from("a").localField("b").foreignField("c").as(null);
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void lookupBuilderBuildsCorrectClause() {
+
+		LookupOperation lookupOperation = LookupOperation.newLookup().from("a").localField("b").foreignField("c").as("d");
+
+		DBObject lookupClause = extractDbObjectFromLookupOperation(lookupOperation);
+
+		assertThat(lookupClause,
+				isBsonObject().containing("from", "a") //
+						.containing("localField", "b") //
+						.containing("foreignField", "c") //
+						.containing("as", "d"));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void lookupBuilderExposesFields() {
+
+		LookupOperation lookupOperation = LookupOperation.newLookup().from("a").localField("b").foreignField("c").as("d");
+
+		assertThat(lookupOperation.getFields().exposesNoFields(), is(false));
+		assertThat(lookupOperation.getFields().exposesSingleFieldOnly(), is(true));
+		assertThat(lookupOperation.getFields().getField("d"), notNullValue());
+	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/TypeBasedAggregationOperationContextUnitTests.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.
@@ -32,6 +32,7 @@ import org.springframework.core.convert.converter.Converter;
 import org.springframework.core.convert.support.GenericConversionService;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.annotation.PersistenceConstructor;
+import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mapping.model.MappingException;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.ExposedField;
 import org.springframework.data.mongodb.core.aggregation.ExposedFields.FieldReference;
@@ -51,6 +52,7 @@ import com.mongodb.DBObject;
  * 
  * @author Oliver Gierke
  * @author Thomas Darimont
+ * @author Mark Paluch
  */
 @RunWith(MockitoJUnitRunner.class)
 public class TypeBasedAggregationOperationContextUnitTests {
@@ -187,6 +189,96 @@ public class TypeBasedAggregationOperationContextUnitTests {
 		assertThat(definition.get("_id"), is(equalTo((Object) "$counter_name")));
 	}
 
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void lookupShouldInheritFieldsFromInheritingAggregationOperation() {
+
+		TypeBasedAggregationOperationContext context = getContext(MeterData.class);
+		TypedAggregation<MeterData> agg = newAggregation(MeterData.class,
+				lookup("OtherCollection", "resourceId", "otherId", "lookup"), sort(Direction.ASC, "resourceId"));
+
+		DBObject dbo = agg.toDbObject("meterData", context);
+		DBObject sort = getPipelineElementFromAggregationAt(dbo, 1);
+
+		DBObject definition = (DBObject) sort.get("$sort");
+
+		assertThat(definition.get("resourceId"), is(equalTo((Object) 1)));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void groupLookupShouldInheritFieldsFromPreviousAggregationOperation() {
+
+		TypeBasedAggregationOperationContext context = getContext(MeterData.class);
+		TypedAggregation<MeterData> agg = newAggregation(MeterData.class, group().min("resourceId").as("foreignKey"),
+				lookup("OtherCollection", "foreignKey", "otherId", "lookup"), sort(Direction.ASC, "foreignKey"));
+
+		DBObject dbo = agg.toDbObject("meterData", context);
+		DBObject sort = getPipelineElementFromAggregationAt(dbo, 2);
+
+		DBObject definition = (DBObject) sort.get("$sort");
+
+		assertThat(definition.get("foreignKey"), is(equalTo((Object) 1)));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void lookupGroupAggregationShouldUseCorrectGroupField() {
+
+		TypeBasedAggregationOperationContext context = getContext(MeterData.class);
+		TypedAggregation<MeterData> agg = newAggregation(MeterData.class,
+				lookup("OtherCollection", "resourceId", "otherId", "lookup"),
+				group().min("lookup.otherkey").as("something_totally_different"));
+
+		DBObject dbo = agg.toDbObject("meterData", context);
+		DBObject group = getPipelineElementFromAggregationAt(dbo, 1);
+
+		DBObject definition = (DBObject) group.get("$group");
+		DBObject field = (DBObject) definition.get("something_totally_different");
+
+		assertThat(field.get("$min"), is(equalTo((Object) "$lookup.otherkey")));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test
+	public void lookupGroupAggregationShouldOverwriteExposedFields() {
+
+		TypeBasedAggregationOperationContext context = getContext(MeterData.class);
+		TypedAggregation<MeterData> agg = newAggregation(MeterData.class,
+				lookup("OtherCollection", "resourceId", "otherId", "lookup"),
+				group().min("lookup.otherkey").as("something_totally_different"),
+				sort(Direction.ASC, "something_totally_different"));
+
+		DBObject dbo = agg.toDbObject("meterData", context);
+		DBObject sort = getPipelineElementFromAggregationAt(dbo, 2);
+
+		DBObject definition = (DBObject) sort.get("$sort");
+
+		assertThat(definition.get("something_totally_different"), is(equalTo((Object) 1)));
+	}
+
+	/**
+	 * @see DATAMONGO-1326
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void lookupGroupAggregationShouldFailInvalidFieldReference() {
+
+		TypeBasedAggregationOperationContext context = getContext(MeterData.class);
+		TypedAggregation<MeterData> agg = newAggregation(MeterData.class,
+				lookup("OtherCollection", "resourceId", "otherId", "lookup"),
+				group().min("lookup.otherkey").as("something_totally_different"), sort(Direction.ASC, "resourceId"));
+
+		agg.toDbObject("meterData", context);
+	}
+
 	@Document(collection = "person")
 	public static class FooPerson {
 

src/main/asciidoc/reference/mongodb.adoc

@@ -1634,6 +1634,7 @@ The MongoDB Aggregation Framework provides the following types of Aggregation Op
 * String Aggregation Operators
 * Date Aggregation Operators
 * Conditional Aggregation Operators
+* Lookup Aggregation Operators
 
 At the time of this writing we provide support for the following Aggregation Operations in Spring Data MongoDB.