DATAMONGO-979 - Add support for $size expression in project and group aggregation pipeline.

Introduced AggregationExpression interface to be able to represent arbitrary MongoDB expressions that can be used in projection and group operations. Supported function expressions are provided via the AggregationFunctionExpressions enum.

Original pull request: #272.

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2015 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 com.mongodb.DBObject;
+
+/**
+ * An {@link AggregationExpression} can be used with field expressions in aggregation pipeline stages like {@code project} and
+ * {@code group}.
+ * 
+ * @author Thomas Darimont
+ */
+public interface AggregationExpression {
+
+	/**
+	 * @param context
+	 * @return
+	 */
+	DBObject toDbObject(AggregationOperationContext context);
+}

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

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2015 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 java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.util.Assert;
+
+import com.mongodb.BasicDBObject;
+import com.mongodb.DBObject;
+
+/**
+ * An enum of supported {@link AggregationExpression}s in aggregation pipeline stages.
+ * 
+ * @author Thomas Darimont
+ * @since 1.10
+ */
+public enum AggregationFunctionExpressions {
+
+	SIZE;
+
+	/**
+	 * Returns an {@link AggregationExpression} build from the current {@link Enum} name and the given {@code params}.
+	 * 
+	 * @param params must not be {@literal null}
+	 * @return
+	 */
+	public AggregationExpression of(Object... params) {
+
+		Assert.notNull(params, "Params must not be null!");
+
+		return new FunctionExpression(name().toLowerCase(), params);
+	}
+
+	/**
+	 * An {@link AggregationExpression} representing a function call.
+	 * 
+	 * @author Thomas Darimont
+	 * @since 1.10
+	 */
+	static class FunctionExpression implements AggregationExpression {
+
+		private final String name;
+		private final Object[] values;
+
+		public FunctionExpression(String name, Object[] values) {
+
+			Assert.hasText(name, "Name must not be null!");
+			Assert.notNull(values, "Values must not be null!");
+
+			this.name = name;
+			this.values = values;
+		}
+
+		/* (non-Javadoc)
+		 * @see org.springframework.data.mongodb.core.aggregation.Expression#toDbObject(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)
+		 */
+		@Override
+		public DBObject toDbObject(AggregationOperationContext context) {
+
+			List<Object> args = new ArrayList<Object>(values.length);
+			for (int i = 0; i < values.length; i++) {
+				args.add(unpack(values[i], context));
+			}
+
+			return new BasicDBObject("$" + name, args);
+		}
+
+		private Object unpack(Object value, AggregationOperationContext context) {
+
+			if (value instanceof AggregationExpression) {
+				return ((AggregationExpression) value).toDbObject(context);
+			}
+
+			if (value instanceof Field) {
+				return context.getReference((Field) value).toString();
+			}
+
+			return value;
+		}
+	}
+}

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

@@ -193,6 +193,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		return newBuilder(GroupOps.LAST, reference, null);
 	}
 
+	/**
+	 * Generates an {@link GroupOperationBuilder} for an {@code $last}-expression for the given {@link AggregationExpression}.
+	 * 
+	 * @param expr
+	 * @return
+	 */
+	public GroupOperationBuilder last(AggregationExpression expr) {
+		return newBuilder(GroupOps.LAST, null, expr);
+	}
+
 	/**
 	 * Generates an {@link GroupOperationBuilder} for a {@code $first}-expression for the given field-reference.
 	 * 
@@ -203,6 +213,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		return newBuilder(GroupOps.FIRST, reference, null);
 	}
 
+	/**
+	 * Generates an {@link GroupOperationBuilder} for a {@code $first}-expression for the given {@link AggregationExpression}.
+	 * 
+	 * @param expr
+	 * @return
+	 */
+	public GroupOperationBuilder first(AggregationExpression expr) {
+		return newBuilder(GroupOps.FIRST, null, expr);
+	}
+
 	/**
 	 * Generates an {@link GroupOperationBuilder} for an {@code $avg}-expression for the given field-reference.
 	 * 
@@ -213,6 +233,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		return newBuilder(GroupOps.AVG, reference, null);
 	}
 
+	/**
+	 * Generates an {@link GroupOperationBuilder} for an {@code $avg}-expression for the given {@link AggregationExpression}.
+	 * 
+	 * @param expr
+	 * @return
+	 */
+	public GroupOperationBuilder avg(AggregationExpression expr) {
+		return newBuilder(GroupOps.AVG, null, expr);
+	}
+
 	/**
 	 * Generates an {@link GroupOperationBuilder} for an {@code $push}-expression for the given field-reference.
 	 * 
@@ -247,6 +277,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		return newBuilder(GroupOps.MIN, reference, null);
 	}
 
+	/**
+	 * Generates an {@link GroupOperationBuilder} for an {@code $min}-expression that for the given {@link AggregationExpression}.
+	 * 
+	 * @param expr
+	 * @return
+	 */
+	public GroupOperationBuilder min(AggregationExpression expr) {
+		return newBuilder(GroupOps.MIN, null, expr);
+	}
+
 	/**
 	 * Generates an {@link GroupOperationBuilder} for an {@code $max}-expression that for the given field-reference.
 	 * 
@@ -257,6 +297,16 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		return newBuilder(GroupOps.MAX, reference, null);
 	}
 
+	/**
+	 * Generates an {@link GroupOperationBuilder} for an {@code $max}-expression that for the given {@link AggregationExpression}.
+	 * 
+	 * @param expr
+	 * @return
+	 */
+	public GroupOperationBuilder max(AggregationExpression expr) {
+		return newBuilder(GroupOps.MAX, null, expr);
+	}
+
 	private GroupOperationBuilder newBuilder(Keyword keyword, String reference, Object value) {
 		return new GroupOperationBuilder(this, new Operation(keyword, null, reference, value));
 	}
@@ -369,6 +419,11 @@ public class GroupOperation implements FieldsExposingAggregationOperation {
 		public Object getValue(AggregationOperationContext context) {
 
 			if (reference == null) {
+
+				if (value instanceof AggregationExpression) {
+					return ((AggregationExpression) value).toDbObject(context);
+				}
+
 				return value;
 			}
 

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

@@ -121,6 +121,10 @@ public class ProjectionOperation implements FieldsExposingAggregationOperation {
 		return new ExpressionProjectionOperationBuilder(expression, this, params);
 	}
 
+	public ProjectionOperationBuilder and(AggregationExpression expression) {
+		return new ProjectionOperationBuilder(expression, this, null);
+	}
+
 	/**
 	 * Excludes the given fields from the projection.
 	 * 
@@ -420,9 +424,13 @@ public class ProjectionOperation implements FieldsExposingAggregationOperation {
 
 			if (this.previousProjection != null) {
 				return this.operation.andReplaceLastOneWith(this.previousProjection.withAlias(alias));
-			} else {
-				return this.operation.and(new FieldProjection(Fields.field(alias, name), null));
 			}
+
+			if (value instanceof AggregationExpression) {
+				return this.operation.and(new ExpressionProjection(Fields.field(alias), (AggregationExpression) value));
+			}
+
+			return this.operation.and(new FieldProjection(Fields.field(alias, name), null));
 		}
 
 		/**
@@ -552,6 +560,10 @@ public class ProjectionOperation implements FieldsExposingAggregationOperation {
 			return project("mod", Fields.field(fieldReference));
 		}
 
+		public ProjectionOperationBuilder size() {
+			return project("size");
+		}
+
 		/* 
 		 * (non-Javadoc)
 		 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#toDBObject(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)
@@ -940,4 +952,31 @@ public class ProjectionOperation implements FieldsExposingAggregationOperation {
 		 */
 		public abstract DBObject toDBObject(AggregationOperationContext context);
 	}
+
+	/**
+	 * @author Thomas Darimont
+	 */
+	static class ExpressionProjection extends Projection {
+
+		private final AggregationExpression expression;
+		private final Field field;
+
+		/**
+		 * Creates a new {@link ExpressionProjection}.
+		 * 
+		 * @param field
+		 * @param expression
+		 */
+		public ExpressionProjection(Field field, AggregationExpression expression) {
+
+			super(field);
+			this.field = field;
+			this.expression = expression;
+		}
+
+		@Override
+		public DBObject toDBObject(AggregationOperationContext context) {
+			return new BasicDBObject(field.getName(), expression.toDbObject(context));
+		}
+	}
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2015 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.
@@ -17,8 +17,11 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.*;
+import static org.springframework.data.mongodb.core.aggregation.AggregationFunctionExpressions.*;
 import static org.springframework.data.mongodb.core.aggregation.Fields.*;
 
+import java.util.Arrays;
+
 import org.junit.Test;
 import org.springframework.data.mongodb.core.DBObjectTestUtils;
 
@@ -29,6 +32,7 @@ import com.mongodb.DBObject;
  * Unit tests for {@link GroupOperation}.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 public class GroupOperationUnitTests {
 
@@ -183,6 +187,23 @@ public class GroupOperationUnitTests {
 		assertThat(push, is((DBObject) new BasicDBObject("$addToSet", 42)));
 	}
 
+	/**
+	 * @see DATAMONGO-979
+	 */
+	@Test
+	public void shouldRenderSizeExpressionInGroup() {
+
+		GroupOperation groupOperation = Aggregation //
+				.group("username") //
+				.first(SIZE.of(field("tags"))) //
+				.as("tags_count");
+
+		DBObject groupClause = extractDbObjectFromGroupOperation(groupOperation);
+		DBObject tagsCount = DBObjectTestUtils.getAsDBObject(groupClause, "tags_count");
+
+		assertThat(tagsCount.get("$first"), is((Object) new BasicDBObject("$size", Arrays.asList("$tags"))));
+	}
+
 	private DBObject extractDbObjectFromGroupOperation(GroupOperation groupOperation) {
 		DBObject dbObject = groupOperation.toDBObject(Aggregation.DEFAULT_CONTEXT);
 		DBObject groupClause = DBObjectTestUtils.getAsDBObject(dbObject, "$group");

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013 the original author or authors.
+ * Copyright 2013-2015 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.
@@ -17,6 +17,8 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
+import static org.springframework.data.mongodb.core.aggregation.AggregationFunctionExpressions.*;
+import static org.springframework.data.mongodb.core.aggregation.Fields.*;
 import static org.springframework.data.mongodb.util.DBObjectUtils.*;
 
 import java.util.Arrays;
@@ -334,6 +336,41 @@ public class ProjectionOperationUnitTests {
 						"$date", 86400000))))));
 	}
 
+	/**
+	 * @see DATAMONGO-979
+	 */
+	@Test
+	public void shouldRenderSizeExpressionInProjection() {
+
+		ProjectionOperation operation = Aggregation //
+				.project() //
+				.and("tags") //
+				.size()//
+				.as("tags_count");
+
+		DBObject dbObject = operation.toDBObject(Aggregation.DEFAULT_CONTEXT);
+
+		DBObject projected = exctractOperation("$project", dbObject);
+		assertThat(projected.get("tags_count"), is((Object) new BasicDBObject("$size", Arrays.asList("$tags"))));
+	}
+
+	/**
+	 * @see DATAMONGO-979
+	 */
+	@Test
+	public void shouldRenderGenericSizeExpressionInProjection() {
+
+		ProjectionOperation operation = Aggregation //
+				.project() //
+				.and(SIZE.of(field("tags"))) //
+				.as("tags_count");
+
+		DBObject dbObject = operation.toDBObject(Aggregation.DEFAULT_CONTEXT);
+
+		DBObject projected = exctractOperation("$project", dbObject);
+		assertThat(projected.get("tags_count"), is((Object) new BasicDBObject("$size", Arrays.asList("$tags"))));
+	}
+
 	private static DBObject exctractOperation(String field, DBObject fromProjectClause) {
 		return (DBObject) fromProjectClause.get(field);
 	}