support for $add and $substract in projection

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

@@ -15,6 +15,12 @@
  */
 package org.springframework.data.mongodb.core.aggregation;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EmptyStackException;
+import java.util.List;
+import java.util.Stack;
+
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.mongodb.core.query.Field;
 import org.springframework.util.Assert;
@@ -26,7 +32,7 @@
  * Projection of field to be used in an {@link AggregationPipeline}.
  * <p/>
  * A projection is similar to a {@link Field} inclusion/exclusion but more powerful. It can generate new fields, change
- * values of given field etc.
+ * values of given field etc. 
  * 
  * @author Tobias Trelle
  */
@@ -36,10 +42,10 @@ public class Projection {
 
 private DBObject document = new BasicDBObject();
 
-/** Key of the current field. */
-private String ref;
-
-private String modifier;
+private DBObject rightHandExpression;
+
+/** Stack of key names. Size is 0 or 1. */
+private Stack<String> reference = new Stack<String>();
 
 /** Create an empty projection. */
 public Projection() {
@@ -75,47 +81,57 @@ public final void exclude(String key) {
  */
 public final Projection include(String key) {
 Assert.notNull(key, "Missing key");
-if (canPop()) {
-document.put(pop(), 1);
-}
-push(key);
+
+safePop();
+reference.push(key);
 
 return this;
 }
 
 /**
  * Sets the key for a computed field.
  * 
-v */
+ */
 public final Projection as(String key) {
 Assert.notNull(key, "Missing key");
 
-document.put( key, safeReference(pop()) );
+try {
+document.put(key, rightHandSide(safeReference(reference.pop())) );
+} catch (EmptyStackException e) {
+throw new InvalidDataAccessApiUsageException("Invalid use of as()", e);
+}
 return this;
 }
 
-private void push(String r) {
-if (ref != null) {
-throw new InvalidDataAccessApiUsageException("No field selected");
-}
-ref = r;
+public final Projection plus(Number n) {
+return arithmeticOperation("add", n);
 }
-
-private String pop() {
-if (ref == null) {
-throw new InvalidDataAccessApiUsageException("No field selected");
-}
-String r = ref;
-ref = null;
-modifier = null;
-
-return r;
+
+public final Projection minus(Number n) {
+return arithmeticOperation("substract", n);
+}
+
+private Projection arithmeticOperation(String op, Number n) {
+Assert.notNull(n, "Missing number");
+
+rightHandExpression = createArrayObject(op, safeReference(reference.peek()), n);
+
+return this;
 }
-
-private boolean canPop() {
-return ref != null;
+
+private DBObject createArrayObject(String op, Object... items) {
+List<Object> list = new ArrayList<Object>();
+Collections.addAll(list, items);
+
+return new BasicDBObject( safeReference(op), list );
 }
-
+
+private void safePop() {
+if ( !reference.empty() ) {
+document.put( reference.pop(), rightHandSide(1) );
+}
+}
+
 private String safeReference(String key) {
 Assert.notNull(key);
 
@@ -125,11 +141,15 @@ private String safeReference(String key) {
 return key;
 }
 }
+
+private Object rightHandSide(Object defaultValue) {
+Object value = rightHandExpression != null ? rightHandExpression : defaultValue;
+rightHandExpression = null;
+return value;
+}
 
 DBObject toDBObject() {
-if (canPop()) {
-document.put(pop(), 1);
-}
+safePop();
 return document;
 }
 

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

@@ -4,6 +4,8 @@
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.junit.Assert.assertThat;
 
+import java.util.List;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
@@ -19,94 +21,143 @@ public class ProjectionTests {
 
 /** Unit under test. */
 private Projection projection;
-
-@Before public void setUp() {
+
+@Before
+public void setUp() {
 projection = new Projection();
 }
-
-@Test public void emptyProjection() {
+
+@Test
+public void emptyProjection() {
 // when
 DBObject raw = projection.toDBObject();
-
+
 // then
-assertThat( raw, notNullValue() );
-assertThat( raw.toMap().isEmpty(), is(true) );
+assertThat(raw, notNullValue());
+assertThat(raw.toMap().isEmpty(), is(true));
 }
-
+
 @Test(expected = IllegalArgumentException.class)
 public void shouldDetectNullIncludesInConstructor() {
 // when
-new Projection((String[])null);
+new Projection((String[]) null);
 // then: throw expected exception
 }
-
-@Test public void includesWithConstructor() {
+
+@Test
+public void includesWithConstructor() {
 // given
 projection = new Projection("a", "b");
-
+
 // when
 DBObject raw = projection.toDBObject();
-
+
 // then
-assertThat( raw, notNullValue() );
-assertThat( raw.toMap().size(), is(3) );
-assertThat( (Integer)raw.get("_id"), is(0) );
-assertThat( (Integer)raw.get("a"), is(1) );
-assertThat( (Integer)raw.get("b"), is(1) );
+assertThat(raw, notNullValue());
+assertThat(raw.toMap().size(), is(3));
+assertThat((Integer) raw.get("_id"), is(0));
+assertThat((Integer) raw.get("a"), is(1));
+assertThat((Integer) raw.get("b"), is(1));
 }
-
-@Test public void include() {
+
+@Test
+public void include() {
 // given
 projection.include("a");
-
+
 // when
 DBObject raw = projection.toDBObject();
-
+
 // then
 assertSingleDBObject("a", 1, raw);
 }
 
-@Test public void exclude() {
+@Test
+public void exclude() {
 // given
 projection.exclude("a");
-
+
 // when
 DBObject raw = projection.toDBObject();
-
+
 // then
 assertSingleDBObject("a", 0, raw);
 }
-
-@Test public void includeAlias() {
+
+@Test
+public void includeAlias() {
 // given
 projection.include("a").as("b");
-
+
 // when
 DBObject raw = projection.toDBObject();
-
+
 // then
 assertSingleDBObject("b", "$a", raw);
 }
 
-@Test(expected = InvalidDataAccessApiUsageException.class ) 
+@Test(expected = InvalidDataAccessApiUsageException.class)
 public void shouldDetectAliasWithoutInclude() {
 // when
 projection.as("b");
 // then: throw expected exception
 }
 
-@Test(expected = InvalidDataAccessApiUsageException.class ) 
+@Test(expected = InvalidDataAccessApiUsageException.class)
 public void shouldDetectDuplicateAlias() {
 // when
 projection.include("a").as("b").as("c");
 // then: throw expected exception
 }
+
+@Test
+public void plus() {
+// given
+projection.include("a").plus(10);
+
+// when
+DBObject raw = projection.toDBObject();
+
+// then
+assertNotNullDBObject(raw);
+DBObject addition = (DBObject)raw.get("a");
+assertNotNullDBObject(addition);
+@SuppressWarnings("unchecked")
+List<Object> summands = (List<Object>)addition.get("$add");
+assertThat( summands, notNullValue() );
+assertThat( summands.size(), is(2) );
+assertThat( (String)summands.get(0), is("$a") );
+assertThat( (Integer)summands.get(1), is (10) );
+}
+
+@Test
+public void plusWithAlias() {
+// given
+projection.include("a").plus(10).as("b");
+
+// when
+DBObject raw = projection.toDBObject();
+
+// then
+assertNotNullDBObject(raw);
+DBObject addition = (DBObject)raw.get("b");
+assertNotNullDBObject(addition);
+@SuppressWarnings("unchecked")
+List<Object> summands = (List<Object>)addition.get("$add");
+assertThat( summands, notNullValue() );
+assertThat( summands.size(), is(2) );
+assertThat( (String)summands.get(0), is("$a") );
+assertThat( (Integer)summands.get(1), is (10) );
+}
 
 
 private static void assertSingleDBObject(String key, Object value, DBObject doc) {
-assertThat( doc, notNullValue() );
-assertThat( doc.toMap().size(), is(1) );
-assertThat( doc.get(key), is(value) );
-}
-
+assertNotNullDBObject(doc);
+assertThat(doc.get(key), is(value));
+}
+
+private static void assertNotNullDBObject(DBObject doc) {
+assertThat(doc, notNullValue());
+assertThat(doc.toMap().size(), is(1));
+}
 }