SERVER-14466 test use of DBRef fields in aggregation and find() expressions

4 files changed, 254 insertions, 2 deletions

diff --git a/jstests/aggregation/dbref.js b/jstests/aggregation/dbref.js
new file mode 100644
index 00000000000..32e6ce428aa
--- /dev/null
+++ b/jstests/aggregation/dbref.js
@@ -0,0 +1,200 @@
+/**
+ * Check that the special $-prefixed field names $ref, $id and $db all work in expressions, $group,
+ * and $lookup.
+ *
+ * Uses $lookup, which doesn't support sharded foreign collection.
+ * DBRef fields aren't supported in agg pre 4.4.
+ * @tags: [assumes_unsharded_collection, requires_fcv_44]
+ */
+(function() {
+const coll = db.dbref_in_expression;
+const otherColl = db.dbref_in_expression_2;
+
+coll.drop();
+otherColl.drop();
+
+assert.commandWorked(otherColl.insert({_id: "id0", x: 1}));
+assert.commandWorked(otherColl.insert({_id: "id1", x: 2}));
+
+assert.commandWorked(coll.insert({
+    _id: 0,
+    link: new DBRef(otherColl.getName(), "id0", db.getName()),
+    linkArray: [
+        new DBRef(otherColl.getName(), "id0", db.getName()),
+        new DBRef(otherColl.getName(), "id1", db.getName())
+    ]
+}));
+
+function projectOnlyPipeline(projection) {
+    const aggRes = coll.aggregate({$project: projection}).toArray();
+    const findRes = coll.find({}, projection).toArray();
+    assert.sameMembers(findRes, aggRes);
+    return aggRes;
+}
+
+// Refer to a DBRef sub-field in a projection.
+assert.eq(projectOnlyPipeline({refVal: "$link.$ref"}), [{_id: 0, refVal: otherColl.getName()}]);
+assert.eq(projectOnlyPipeline({refVal: "$linkArray.$ref"}),
+          [{_id: 0, refVal: [otherColl.getName(), otherColl.getName()]}]);
+
+assert.eq(projectOnlyPipeline({idVal: "$link.$id"}), [{_id: 0, idVal: "id0"}]);
+assert.eq(projectOnlyPipeline({idVal: "$linkArray.$id"}), [{_id: 0, idVal: ["id0", "id1"]}]);
+
+assert.eq(projectOnlyPipeline({idVal: "$link.$db"}), [{_id: 0, idVal: db.getName()}]);
+assert.eq(projectOnlyPipeline({idVal: "$linkArray.$db"}),
+          [{_id: 0, idVal: [db.getName(), db.getName()]}]);
+
+// Use a DBRef sub-field in an expression.
+assert.eq(projectOnlyPipeline({idLen: {$strLenCP: "$link.$id"}}), [{_id: 0, idLen: "id0".length}]);
+
+// Project away DBRef values.
+assert.eq(projectOnlyPipeline({link: {$ref: 0}, linkArray: 0}),
+          [{_id: 0, link: {$id: "id0", $db: db.getName()}}]);
+
+assert.eq(projectOnlyPipeline({link: 0, linkArray: {$id: 0}}), [{
+              _id: 0,
+              linkArray: [
+                  {$ref: otherColl.getName(), $db: db.getName()},
+                  {$ref: otherColl.getName(), $db: db.getName()}
+              ]
+          }]);
+
+// Assigning to a DBRef field.
+assert.eq(projectOnlyPipeline({link: {$ref: 1, $id: 1, $db: "someOtherDB"}}),
+          [{_id: 0, link: new DBRef(otherColl.getName(), "id0", "someOtherDB")}]);
+
+// While not a 'feature' we advertise, it is allowed to assign to top-level DBRef fields.
+assert.eq(projectOnlyPipeline({$ref: "$link.$ref"}), [{_id: 0, $ref: otherColl.getName()}]);
+
+// One cannot refer to a top-level DBRef field, however, as it will be interpreted as a variable
+// dereference.
+const err = assert.throws(() => coll.aggregate({$project: {x: "$$ref"}}).toArray());
+assert.eq(err.code, 17276);
+
+// It can be accessed through $$ROOT, however.
+assert.eq(coll.aggregate([
+                  // Rather than go through the trouble of inserting a document with a top-level
+                  // $-prefixed field, create one in an intermediate $project stage.
+                  {$project: {"$ref": "hello world"}},
+                  // Make sure that no optimization coalesces the above projection stage with the
+                  // below one.
+                  {$_internalInhibitOptimization: {}},
+                  {$project: {x: "$$ROOT.$ref"}}
+              ])
+              .toArray(),
+          [{_id: 0, x: "hello world"}]);
+
+// Do a count (using $group) on a DBRef field.
+assert.eq(coll.aggregate({$group: {_id: "$link.$db", count: {$sum: 1}}}).toArray(),
+          [{_id: db.getName(), count: 1}]);
+
+// Refer to a DBRef field in an accumulator.
+assert.eq(coll.aggregate({$group: {_id: "$link.$db", count: {$sum: {$size: "$linkArray.$ref"}}}})
+              .toArray(),
+          [{_id: db.getName(), count: 2}]);
+
+// Use $lookup with a DBRef.
+
+// Equality match version.
+const lookupEqualityPipeline = [{$lookup: {from: otherColl.getName(),
+                                           localField: "link.$id",
+                                           foreignField: "_id",
+                                           as: "joinedField"}},
+                        {$project: {link: 0, linkArray: 0}}];
+assert.eq(coll.aggregate(lookupEqualityPipeline).toArray(),
+          [{_id: 0, joinedField: [{_id: "id0", x: 1}]}]);
+
+// Foreign pipeline.
+const lookupSubPipePipeline = [{$lookup: {from: otherColl.getName(),
+                                          let: {idsWanted: "$linkArray.$id"},
+                                          pipeline: [{$match: {$expr: {$in: ["$_id", "$$idsWanted"]}}}],
+                                          as: "joinedField"}},
+                        {$project: {link: 0, linkArray: 0}}];
+assert.eq(coll.aggregate(lookupSubPipePipeline).toArray(),
+          [{_id: 0, joinedField: [{_id: "id0", x: 1}, {_id: "id1", x: 2}]}]);
+
+(function testGraphLookup() {
+    // $graphLookup using DBRef.
+    const graphLookupColl = db.dbref_graph_lookup;
+    graphLookupColl.drop();
+
+    // id0 -> id1 -> id2 -> id0
+    assert.commandWorked(graphLookupColl.insert(
+        {_id: "id0", link: new DBRef(graphLookupColl.getName(), "id1", db.getName())}));
+    assert.commandWorked(graphLookupColl.insert(
+        {_id: "id1", link: new DBRef(graphLookupColl.getName(), "id2", db.getName())}));
+    assert.commandWorked(graphLookupColl.insert(
+        {_id: "id2", link: new DBRef(graphLookupColl.getName(), "id0", db.getName())}));
+
+    // id3 -> id4
+    assert.commandWorked(graphLookupColl.insert(
+        {_id: "id3", link: new DBRef(graphLookupColl.getName(), "id4", db.getName())}));
+    assert.commandWorked(graphLookupColl.insert({_id: "id4", link: null}));
+
+    const graphLookupPipeline = [{
+        $graphLookup: {
+            from: graphLookupColl.getName(),
+            startWith: "$link.$id",
+            connectFromField: "link.$id",
+            connectToField: "_id",
+            as: "connectedDocuments"
+        }
+    },
+                                 {$sort: {_id: 1}}];
+
+    const res = graphLookupColl.aggregate(graphLookupPipeline).toArray();
+    // id0, id1, and id2 are all connected.
+    assert.eq(res[0].connectedDocuments.length, 3);
+    assert.eq(res[1].connectedDocuments.length, 3);
+    assert.eq(res[2].connectedDocuments.length, 3);
+
+    // id3 is connected to id4.
+    assert.eq(res[3].connectedDocuments.length, 1);
+
+    // id4 is connected to nothing.
+    assert.eq(res[4].connectedDocuments.length, 0);
+    assert.eq(res.length, 5);
+})();
+
+// Distinct command for a dbref field.
+assert.eq(coll.distinct("link.$ref"), [otherColl.getName()]);
+assert.eq(coll.distinct("link.$id"), ["id0"]);
+assert.eq(coll.distinct("link.$db"), [db.getName()]);
+
+// $merge pipeline.
+const thirdColl = db.dbref_in_expression_3;
+thirdColl.drop();
+assert.commandWorked(thirdColl.insert({_id: 0, a: 1}));
+assert.commandWorked(coll.createIndex({"link.$ref": 1}, {unique: true}));
+
+// Merge a document with a 'sentinel' field into the original collection using 'link.$ref' as the
+// "on" field.
+thirdColl
+    .aggregate([
+        {$project: {"link.$ref": otherColl.getName(), "link.$id": "id0", sentinel: "foo"}},
+        {$merge: {into: coll.getName(), on: "link.$ref", whenMatched: "replace"}}
+    ])
+    .itcount();
+
+// Check that the merge worked.
+assert.eq(coll.find({sentinel: "foo"}).itcount(), 1);
+
+// Merge using an update pipeline.
+thirdColl
+    .aggregate([
+        {$project: {"link.$ref": otherColl.getName(), "link.$id": "id0", sentinel: "foo"}},
+        {
+            $merge: {
+                into: coll.getName(),
+                on: "link.$ref",
+                whenMatched: [{
+                    $project:
+                        {"link.$ref": "otherRef", "link.$id": "otherId", "link.$db": "otherDB"}
+                }],
+                whenNotMatched: "discard"
+            }
+        }
+    ])
+    .itcount();
+assert.eq(coll.find().toArray()[0], {_id: 0, link: new DBRef("otherRef", "otherId", "otherDB")});
+})();
diff --git a/jstests/core/elemMatchProjection.js b/jstests/core/elemMatchProjection.js
index 1e6632ea2a5..1bf4fe11217 100644
--- a/jstests/core/elemMatchProjection.js
+++ b/jstests/core/elemMatchProjection.js
@@ -1,5 +1,4 @@
-// @tags: [requires_getmore]
-
+// @tags: [requires_getmore, requires_fcv_44]
 // Tests for $elemMatch projections and $ positional operator projection.
 (function() {
 "use strict";
@@ -65,6 +64,18 @@ for (let i = 0; i < 100; i++) {
     bulk.insert({_id: nextId(), group: 12, x: {y: [{a: 1, b: 1}, {a: 1, b: 2}]}});
     bulk.insert({_id: nextId(), group: 13, x: [{a: 1, b: 1}, {a: 1, b: 2}]});
     bulk.insert({_id: nextId(), group: 13, x: [{a: 1, b: 2}, {a: 1, b: 1}]});
+
+    // Array of DBRefs. Don't actually try to dereference them, though, as they point to
+    // non-existing collections.
+    bulk.insert({
+        _id: nextId(),
+        group: 14,
+        x: [
+            new DBRef("otherCollection", "id0", db.getName()),
+            new DBRef("otherCollection", "id1", db.getName()),
+            new DBRef("otherCollection2", "id2", db.getName())
+        ]
+    });
 }
 assert.commandWorked(bulk.execute());
 
@@ -217,6 +228,16 @@ assert.eq({"x": [{"a": 1, "b": 2}], "y": [{"c": 3, "d": 4}]},
               .toArray()[0],
           "multiple $elemMatch on unique fields 1");
 
+// Perform a $elemMatch on a DBRef field.
+assert.eq(coll.find({group: 14}, {x: {$elemMatch: {$id: "id0"}}}).toArray()[0].x,
+          [new DBRef("otherCollection", "id0", db.getName())]);
+
+assert.eq(coll.find({group: 14}, {x: {$elemMatch: {$ref: "otherCollection2"}}}).toArray()[0].x,
+          [new DBRef("otherCollection2", "id2", db.getName())]);
+
+assert.eq(coll.find({group: 14}, {x: {$elemMatch: {$db: db.getName()}}}).toArray()[0].x,
+          [new DBRef("otherCollection", "id0", db.getName())]);
+
 // Tests involving getMore. Test the $-positional operator across multiple batches.
 let a = coll.find({group: 3, 'x.b': 2}, {'x.$': 1}).sort({_id: 1}).batchSize(1);
 while (a.hasNext()) {
diff --git a/jstests/core/views/dbref_projection.js b/jstests/core/views/dbref_projection.js
new file mode 100644
index 00000000000..963a3cc185e
--- /dev/null
+++ b/jstests/core/views/dbref_projection.js
@@ -0,0 +1,30 @@
+/**
+ * Test projecting DBRef fields ($ref, $id, $db) in views.
+ *
+ * Legacy find() queries do not support views, so must use the find() command.
+ * DBRef fields are not supported in agg pre 4.4.
+ * @tags: [requires_find_command, requires_fcv_44]
+ */
+(function() {
+"use strict";
+
+const viewsDB = db.getSiblingDB("views_dbref_projection");
+assert.commandWorked(viewsDB.dropDatabase());
+
+assert.commandWorked(
+    viewsDB.baseColl.insert({_id: 0, link: new DBRef("otherColl", "someId", viewsDB.getName())}));
+
+assert.commandWorked(viewsDB.runCommand({create: "view", viewOn: "baseColl"}));
+
+// Check that the view and base collection return the same thing.
+function checkViewAndBaseCollection(projection, expectedResult) {
+    const baseRes = viewsDB.baseColl.find({}, projection).toArray();
+    const viewRes = viewsDB.view.find({}, projection).toArray();
+    assert.eq(baseRes, viewRes);
+    assert.eq(expectedResult, baseRes);
+}
+
+checkViewAndBaseCollection({"link.$ref": 1}, [{_id: 0, link: {$ref: "otherColl"}}]);
+checkViewAndBaseCollection({"link.$db": 1}, [{_id: 0, link: {$db: viewsDB.getName()}}]);
+checkViewAndBaseCollection({"link.$id": 1}, [{_id: 0, link: {$id: "someId"}}]);
+}());
diff --git a/jstests/libs/parallelTester.js b/jstests/libs/parallelTester.js
index 525e0a3c802..22b9277042d 100644
--- a/jstests/libs/parallelTester.js
+++ b/jstests/libs/parallelTester.js
@@ -222,6 +222,7 @@ if (typeof _threadInject != "undefined") {
             var requires_find_command = [
                 "update_pipeline_shell_helpers.js",
                 "update_with_pipeline.js",
+                "views/dbref_projection.js",
                 "views/views_aggregation.js",
                 "views/views_change.js",
                 "views/views_drop.js",