implement the $graphLookup pipeline stage

Justin Lee · Justin Lee · commit 83cfd1fb940b · 2016-09-07T15:18:38.000-04:00
resolves JAVA-2187
diff --git a/driver-core/src/main/com/mongodb/client/model/Aggregates.java b/driver-core/src/main/com/mongodb/client/model/Aggregates.java
@@ -113,6 +113,46 @@ public static Bson lookup(final String from, final String localField, final Stri
                                                    .append("as", new BsonString(as)));
     }
 
+    /**
+     * Creates a graphLookup pipeline stage for the specified filter
+     *
+     * @param <TExpression>     the expression type
+     * @param from             the collection to query
+     * @param startWith        the expression to start the graph lookup with
+     * @param connectFromField the from field
+     * @param connectToField   the to field
+     * @param as               name of field in output document
+     * @return the graphLookup pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/graphLookup/ graphLookup
+     * @mongodb.server.release 3.4
+     * @since 3.4
+     */
+    public static <TExpression> Bson graphLookup(final String from, final TExpression startWith, final String connectFromField,
+                                                 final String connectToField, final String as) {
+        return graphLookup(from, startWith, connectFromField, connectToField, as, new GraphLookupOptions());
+    }
+
+    /**
+     * Creates a graphLookup pipeline stage for the specified filter
+     *
+     * @param <TExpression>    the expression type
+     * @param from             the collection to query
+     * @param startWith        the expression to start the graph lookup with
+     * @param connectFromField the from field
+     * @param connectToField   the to field
+     * @param as               name of field in output document
+     * @param options          optional values for the graphLookup
+     * @return the graphLookup pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/graphLookup/ graphLookup
+     * @mongodb.server.release 3.4
+     * @since 3.4
+     */
+    public static <TExpression> Bson graphLookup(final String from, final TExpression startWith, final String connectFromField,
+                                                 final String connectToField, final String as, final GraphLookupOptions options) {
+        notNull("options", options);
+        return new GraphLookupStage<TExpression>(from, startWith, connectFromField, connectToField, as, options);
+    }
+
     /**
      * Creates a $group pipeline stage for the specified filter
      *
@@ -222,6 +262,64 @@ public String toString() {
         }
     }
 
+    private static final class GraphLookupStage<TExpression> implements Bson {
+        private final String from;
+        private final TExpression startWith;
+        private final String connectFromField;
+        private final String connectToField;
+        private final String as;
+        private final GraphLookupOptions options;
+
+        private GraphLookupStage(final String from, final TExpression startWith, final String connectFromField, final String connectToField,
+                                 final String as, final GraphLookupOptions options) {
+            this.from = from;
+            this.startWith = startWith;
+            this.connectFromField = connectFromField;
+            this.connectToField = connectToField;
+            this.as = as;
+            this.options = options;
+        }
+
+        @Override
+        public <TDocument> BsonDocument toBsonDocument(final Class<TDocument> tDocumentClass, final CodecRegistry codecRegistry) {
+            BsonDocumentWriter writer = new BsonDocumentWriter(new BsonDocument());
+
+            writer.writeStartDocument();
+
+            writer.writeStartDocument("$graphLookup");
+
+            writer.writeString("from", from);
+            writer.writeName("startWith");
+            BuildersHelper.encodeValue(writer, startWith, codecRegistry);
+
+            writer.writeString("connectFromField", connectFromField);
+            writer.writeString("connectToField", connectToField);
+            writer.writeString("as", as);
+            if (options.getMaxDepth() != null) {
+                writer.writeInt32("maxDepth", options.getMaxDepth());
+            }
+            if (options.getDepthField() != null) {
+                writer.writeString("depthField", options.getDepthField());
+            }
+
+            writer.writeEndDocument();
+
+            return writer.getDocument();
+        }
+
+        @Override
+        public String toString() {
+            return "GraphLookupStage{"
+                + "as='" + as + '\''
+                + ", connectFromField='" + connectFromField + '\''
+                + ", connectToField='" + connectToField + '\''
+                + ", from='" + from + '\''
+                + ", options=" + options
+                + ", startWith=" + startWith
+                + '}';
+        }
+    }
+
     private static class GroupStage<TExpression> implements Bson {
         private final TExpression id;
         private final List<BsonField> fieldAccumulators;
diff --git a/driver-core/src/main/com/mongodb/client/model/GraphLookupOptions.java b/driver-core/src/main/com/mongodb/client/model/GraphLookupOptions.java
@@ -0,0 +1,87 @@
+/*
+ * Copyright (c) 2008-2016 MongoDB, Inc.
+ *
+ * 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 com.mongodb.client.model;
+
+/**
+ * The options for a graphLookup aggregation pipeline stage
+ *
+ * @mongodb.driver.manual reference/operator/aggregation/graphLookup/ graphLookup
+ * @mongodb.server.release 3.4
+ * @since 3.4
+ */
+public final class GraphLookupOptions {
+    private Integer maxDepth;
+    private String depthField;
+
+    /**
+     * The name of the field in which to store the depth value
+     *
+     * @param field the field name
+     * @return this
+     */
+    public GraphLookupOptions depthField(final String field) {
+        depthField = field;
+        return this;
+    }
+
+    /**
+     * @return the field name
+     */
+    public String getDepthField() {
+        return depthField;
+    }
+
+    /**
+     * Specifies a maximum recursive depth for the $graphLookup.  This number must be non-negative.
+     *
+     * @param max the maximum depth
+     * @return this
+     */
+    public GraphLookupOptions maxDepth(final Integer max) {
+        maxDepth = max;
+        return this;
+    }
+
+    /**
+     * @return the maximum depth
+     */
+    public Integer getMaxDepth() {
+        return maxDepth;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder stringBuilder = new StringBuilder()
+            .append("GraphLookupOptions{");
+        if (depthField != null) {
+            stringBuilder.append("depthField='")
+                         .append(depthField)
+                         .append('\'');
+            if (maxDepth != null) {
+                stringBuilder.append(", ");
+            }
+        }
+        if (maxDepth != null) {
+            stringBuilder
+                .append("maxDepth=")
+                .append(maxDepth);
+        }
+        return stringBuilder
+            .append('}')
+            .toString();
+    }
+}
diff --git a/driver-core/src/test/functional/com/mongodb/client/model/AggregatesFunctionalSpecification.groovy b/driver-core/src/test/functional/com/mongodb/client/model/AggregatesFunctionalSpecification.groovy
@@ -16,6 +16,7 @@ package com.mongodb.client.model
 
 import com.mongodb.MongoNamespace
 import com.mongodb.OperationFunctionalSpecification
+import org.bson.BsonString
 import org.bson.Document
 import org.bson.conversions.Bson
 import spock.lang.IgnoreIf
@@ -32,6 +33,7 @@ import static com.mongodb.client.model.Accumulators.push
 import static com.mongodb.client.model.Accumulators.stdDevPop
 import static com.mongodb.client.model.Accumulators.stdDevSamp
 import static com.mongodb.client.model.Accumulators.sum
+import static com.mongodb.client.model.Aggregates.graphLookup
 import static com.mongodb.client.model.Aggregates.group
 import static com.mongodb.client.model.Aggregates.limit
 import static com.mongodb.client.model.Aggregates.lookup
@@ -231,4 +233,82 @@ class AggregatesFunctionalSpecification extends OperationFunctionalSpecification
         cleanup:
         fromHelper?.drop()
     }
+
+    @IgnoreIf({ !serverVersionAtLeast(asList(3, 3, 9)) })
+    def '$graphLookup'() {
+        given:
+        def fromCollectionName = 'contacts'
+        def fromHelper = getCollectionHelper(new MongoNamespace(getDatabaseName(), fromCollectionName))
+
+        fromHelper.drop()
+
+        fromHelper.insertDocuments(Document.parse('{ _id: 0, name: "Bob Smith", friends: ["Anna Jones", "Chris Green"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 1, name: "Anna Jones", friends: ["Bob Smith", "Chris Green", "Joe Lee"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 2, name: "Chris Green", friends: ["Anna Jones", "Bob Smith"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 3, name: "Joe Lee", friends: ["Anna Jones", "Fred Brown"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 4, name: "Fred Brown", friends: ["Joe Lee"] }'))
+
+        def lookupDoc = graphLookup('contacts', new BsonString('$friends'), 'friends', 'name', 'socialNetwork')
+
+        when:
+        def results = fromHelper.aggregate([lookupDoc])
+
+        then:
+        results[0] ==
+            Document.parse('''{
+                _id: 0,
+                name: "Bob Smith",
+                friends: ["Anna Jones", "Chris Green"],
+                socialNetwork: [
+                    { _id: 3, name: "Joe Lee", friends: ["Anna Jones", "Fred Brown" ] },
+                    { _id: 4, name: "Fred Brown", friends: ["Joe Lee"] },
+                    { _id: 0, name: "Bob Smith", friends: ["Anna Jones", "Chris Green"] },
+                    { _id: 2, name: "Chris Green", friends: ["Anna Jones", "Bob Smith"] },
+                    { _id: 1, name: "Anna Jones", friends: ["Bob Smith", "Chris Green", "Joe Lee"] }
+                ]
+            }''')
+
+        cleanup:
+        fromHelper?.drop()
+    }
+
+    @IgnoreIf({ !serverVersionAtLeast(asList(3, 3, 9)) })
+    def '$graphLookup with options'() {
+        given:
+        def fromCollectionName = 'contacts'
+        def fromHelper = getCollectionHelper(new MongoNamespace(getDatabaseName(), fromCollectionName))
+
+        fromHelper.drop()
+
+        fromHelper.insertDocuments(Document.parse('{ _id: 0, name: "Bob Smith", friends: ["Anna Jones", "Chris Green"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 1, name: "Anna Jones", friends: ["Bob Smith", "Chris Green", "Joe Lee"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 2, name: "Chris Green", friends: ["Anna Jones", "Bob Smith"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 3, name: "Joe Lee", friends: ["Anna Jones", "Fred Brown"] }'))
+        fromHelper.insertDocuments(Document.parse('{ _id: 4, name: "Fred Brown", friends: ["Joe Lee"] }'))
+
+        def lookupDoc = graphLookup('contacts', new BsonString('$friends'), 'friends', 'name', 'socialNetwork',
+                                    new GraphLookupOptions()
+                                        .maxDepth(1)
+                                        .depthField('depth'))
+
+        when:
+        def results = fromHelper.aggregate([lookupDoc])
+
+        then:
+        results[0] ==
+            Document.parse('''{
+                _id: 0,
+                name: "Bob Smith",
+                friends: ["Anna Jones", "Chris Green"],
+                socialNetwork: [
+                    { _id: 3, name: "Joe Lee", friends: ["Anna Jones", "Fred Brown" ], depth:1 },
+                    { _id: 0, name: "Bob Smith", friends: ["Anna Jones", "Chris Green"], depth:1 },
+                    { _id: 2, name: "Chris Green", friends: ["Anna Jones", "Bob Smith"], depth:0 },
+                    { _id: 1, name: "Anna Jones", friends: ["Bob Smith", "Chris Green", "Joe Lee"], depth:0 }
+                ]
+            }''')
+
+        cleanup:
+        fromHelper?.drop()
+    }
 }
diff --git a/driver-core/src/test/unit/com/mongodb/client/model/AggregatesSpecification.groovy b/driver-core/src/test/unit/com/mongodb/client/model/AggregatesSpecification.groovy
@@ -32,6 +32,7 @@ import static com.mongodb.client.model.Accumulators.push
 import static com.mongodb.client.model.Accumulators.stdDevPop
 import static com.mongodb.client.model.Accumulators.stdDevSamp
 import static com.mongodb.client.model.Accumulators.sum
+import static com.mongodb.client.model.Aggregates.graphLookup
 import static com.mongodb.client.model.Aggregates.group
 import static com.mongodb.client.model.Aggregates.limit
 import static com.mongodb.client.model.Aggregates.lookup
@@ -80,6 +81,25 @@ class AggregatesSpecification extends Specification {
             foreignField: "foreignField", as: "as" } }''')
     }
 
+    def 'should render $graphLookup'() {
+        expect:
+        toBson(graphLookup('contacts', '$friends', 'friends', 'name', 'socialNetwork')) ==
+        parse('''{ $graphLookup: { from: "contacts", startWith: "$friends", connectFromField: "friends", connectToField: "name",
+            as: "socialNetwork" } }''')
+        toBson(graphLookup('contacts', '$friends', 'friends', 'name', 'socialNetwork', new GraphLookupOptions().maxDepth(1))) ==
+        parse('''{ $graphLookup: { from: "contacts", startWith: "$friends", connectFromField: "friends", connectToField: "name",
+            as: "socialNetwork", maxDepth: 1 } }''')
+        toBson(graphLookup('contacts', '$friends', 'friends', 'name', 'socialNetwork', new GraphLookupOptions()
+                .maxDepth(1)
+                .depthField('master'))) ==
+        parse('''{ $graphLookup: { from: "contacts", startWith: "$friends", connectFromField: "friends", connectToField: "name",
+            as: "socialNetwork", maxDepth: 1, depthField: "master" } }''')
+        toBson(graphLookup('contacts', '$friends', 'friends', 'name', 'socialNetwork', new GraphLookupOptions()
+                .depthField('master'))) ==
+        parse('''{ $graphLookup: { from: "contacts", startWith: "$friends", connectFromField: "friends", connectToField: "name",
+            as: "socialNetwork", depthField: "master" } }''')
+    }
+
     def 'should render $skip'() {
         expect:
         toBson(skip(5)) == parse('{ $skip : 5 }')
