Implement support for whenMatched pipelines ($merge aggregation) #152

Author: bwaldvogel

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Aggregation.java

@@ -29,7 +29,7 @@
 import de.bwaldvogel.mongo.backend.aggregation.stage.LookupWithPipelineStage;
 import de.bwaldvogel.mongo.backend.aggregation.stage.MatchStage;
 import de.bwaldvogel.mongo.backend.aggregation.stage.MergeStage;
-import de.bwaldvogel.mongo.backend.aggregation.stage.OrderByStage;
+import de.bwaldvogel.mongo.backend.aggregation.stage.SortStage;
 import de.bwaldvogel.mongo.backend.aggregation.stage.OutStage;
 import de.bwaldvogel.mongo.backend.aggregation.stage.ProjectStage;
 import de.bwaldvogel.mongo.backend.aggregation.stage.RedactStage;
@@ -104,7 +104,7 @@ public static Aggregation fromPipeline(List<Document> pipeline, DatabaseResolver
                     break;
                 case "$sort":
                     Document orderBy = (Document) stage.get(stageOperation);
-                    aggregation.addStage(new OrderByStage(orderBy));
+                    aggregation.addStage(new SortStage(orderBy));
                     break;
                 case "$project":
                     Document projection = (Document) stage.get(stageOperation);
@@ -150,7 +150,7 @@ public static Aggregation fromPipeline(List<Document> pipeline, DatabaseResolver
                 case "$sortByCount":
                     Object expression = stage.get(stageOperation);
                     aggregation.addStage(new GroupStage(new Document(ID_FIELD, expression).append("count", new Document("$sum", 1))));
-                    aggregation.addStage(new OrderByStage(new Document("count", -1).append(ID_FIELD, 1)));
+                    aggregation.addStage(new SortStage(new Document("count", -1).append(ID_FIELD, 1)));
                     break;
                 case "$bucket":
                     Document bucket = (Document) stage.get(stageOperation);
@@ -270,6 +270,10 @@ private boolean isLastStage(TerminalStage stage) {
         return stages.indexOf(stage) == stages.size() - 1;
     }
 
+    public List<AggregationStage> getStages() {
+        return stages;
+    }
+
     public boolean isModifying() {
         return stages.stream().anyMatch(AggregationStage::isModifying);
     }

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/Expression.java

@@ -1773,18 +1773,34 @@ static Object evaluate(Object expression, Document document) {
             }
             String value = ((String) expression).substring(1);
             if (value.startsWith("$")) {
-                if (value.equals("$ROOT")) {
-                    return document;
-                } else if (value.startsWith("$ROOT.")) {
-                    String subKey = value.substring("$ROOT.".length());
-                    return Utils.getSubdocumentValue(document, subKey);
-                }
-                Object subdocumentValue = Utils.getSubdocumentValue(document, value);
-                if (!(subdocumentValue instanceof Missing)) {
-                    return subdocumentValue;
-                }
-                String variable = value.substring(1);
-                throw new MongoServerError(17276, "Use of undefined variable: " + variable);
+                final String variableName;
+                if (value.contains(".")) {
+                    variableName = value.substring(0, value.indexOf('.'));
+                } else {
+                    variableName = value;
+                }
+
+                Object variableValue = Utils.getSubdocumentValue(document, variableName);
+                if (variableValue instanceof Missing) {
+                    if (variableName.equals("$ROOT")) {
+                        variableValue = document;
+                    } else {
+                        throw new MongoServerError(17276, "Use of undefined variable: " + variableName.substring(1));
+                    }
+                }
+
+                if (!value.equals(variableName)) {
+                    String path = value.substring(variableName.length() + 1);
+                    final Document evaluatedVariableValue;
+                    if (variableValue instanceof Document) {
+                        evaluatedVariableValue = (Document) variableValue;
+                    } else {
+                        evaluatedVariableValue = (Document) evaluate(variableValue, document);
+                    }
+                    return Utils.getSubdocumentValue(evaluatedVariableValue, path);
+                } else {
+                    return variableValue;
+                }
             }
             return Utils.getSubdocumentValueCollectionAware(document, value);
         } else if (expression instanceof Document) {

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/AbstractLookupStage.java

@@ -9,19 +9,23 @@
 
 abstract class AbstractLookupStage implements AggregationStage {
 
-    protected String stageName = "$lookup";
     static final String FROM = "from";
     static final String AS = "as";
 
+    @Override
+    public String name() {
+        return "$lookup";
+    }
+
     String readStringConfigurationProperty(Document configuration, String name) {
         Object value = configuration.get(name);
         if (value == null) {
-            throw new FailedToParseException("missing '" + name + "' option to " + stageName + " stage specification: " + configuration);
+            throw new FailedToParseException("missing '" + name + "' option to " + name() + " stage specification: " + configuration);
         }
         if (value instanceof String) {
             return (String) value;
         }
-        throw new FailedToParseException("'" + name + "' option to " + stageName + " must be a string, but was type " + Utils.describeType(value));
+        throw new FailedToParseException("'" + name + "' option to " + name() + " must be a string, but was type " + Utils.describeType(value));
     }
 
     Document readOptionalDocumentArgument(Document configuration, String name) {
@@ -32,13 +36,13 @@ Document readOptionalDocumentArgument(Document configuration, String name) {
         if (value instanceof Document) {
             return (Document) value;
         }
-        throw new FailedToParseException(stageName + " argument '" + name + ": " + Json.toJsonValue(value) + "' must be an object, is type " + Utils.describeType(value));
+        throw new FailedToParseException(name() + " argument '" + name + ": " + Json.toJsonValue(value) + "' must be an object, is type " + Utils.describeType(value));
     }
 
     void ensureAllConfigurationPropertiesAreKnown(Document configuration, Set<String> configurationKeys) {
         for (String name : configuration.keySet()) {
             if (!configurationKeys.contains(name)) {
-                String message = "unknown argument to " + stageName + ": " + name;
+                String message = "unknown argument to " + name() + ": " + name;
                 throw new FailedToParseException(message);
             }
         }

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/AddFieldsStage.java

@@ -19,6 +19,11 @@ public AddFieldsStage(Document addFields) {
         this.addFields = addFields;
     }
 
+    @Override
+    public String name() {
+        return "$addFields";
+    }
+
     @Override
     public Stream<Document> apply(Stream<Document> stream) {
         return stream.map(this::projectDocument);

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/AggregationStage.java

@@ -6,6 +6,8 @@
 
 public interface AggregationStage {
 
+    String name();
+
     Stream<Document> apply(Stream<Document> stream);
 
     default boolean isModifying() {

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/BucketStage.java

@@ -37,6 +37,11 @@ public BucketStage(Document document) {
         output = getAndValidateOutput(document.get("output"));
     }
 
+    @Override
+    public String name() {
+        return "$bucket";
+    }
+
     private void validateValuePresent(Object value) {
         if (value == null) {
             throw new MongoServerError(40198, "$bucket requires 'groupBy' and 'boundaries' to be specified.");

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/FacetStage.java

@@ -26,6 +26,11 @@ public FacetStage(Document facetsConfiguration, DatabaseResolver databaseResolve
         }
     }
 
+    @Override
+    public String name() {
+        return "$facet";
+    }
+
     @Override
     public Stream<Document> apply(Stream<Document> stream) {
         List<Document> allDocuments = stream.collect(Collectors.toList());

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/GraphLookupStage.java

@@ -43,7 +43,6 @@ public class GraphLookupStage extends AbstractLookupStage {
     private final MongoCollection<?> collection;
 
     public GraphLookupStage(Document configuration, MongoDatabase mongoDatabase) {
-        stageName = "$graphLookup";
         String from = readStringConfigurationProperty(configuration, FROM);
         collection = mongoDatabase.resolveCollection(from, false);
         readStringConfigurationProperty(configuration, FROM);
@@ -56,6 +55,11 @@ public GraphLookupStage(Document configuration, MongoDatabase mongoDatabase) {
         ensureAllConfigurationPropertiesAreKnown(configuration, CONFIGURATION_KEYS);
     }
 
+    @Override
+    public String name() {
+        return "$graphLookup";
+    }
+
     Integer readOptionalIntegerConfigurationProperty(Document configuration, String name) {
         Object value = configuration.get(name);
         if (value == null) {

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/GroupStage.java

@@ -32,6 +32,11 @@ public GroupStage(Document groupQuery) {
         accumulatorSuppliers = Accumulator.parse(groupQuery);
     }
 
+    @Override
+    public String name() {
+        return "$group";
+    }
+
     @Override
     public Stream<Document> apply(Stream<Document> stream) {
         Map<Object, Collection<Accumulator>> accumulatorsPerKey = new TreeMap<>(ValueComparator.asc());

core/src/main/java/de/bwaldvogel/mongo/backend/aggregation/stage/IndexStatsStage.java

@@ -16,6 +16,11 @@ public IndexStatsStage(MongoCollection<?> collection) {
         this.collection = collection;
     }
 
+    @Override
+    public String name() {
+        return "$indexStats";
+    }
+
     @Override
     public Stream<Document> apply(Stream<Document> stream) {
         return collection.getIndexes().stream()