SERVER-31549 Additional tests for $lookup w/ pipeline syntax

5 files changed, 433 insertions, 64 deletions

diff --git a/jstests/aggregation/bugs/cursor_timeout.js b/jstests/aggregation/bugs/cursor_timeout.js
index 8d61d51378f..817465a0218 100644
--- a/jstests/aggregation/bugs/cursor_timeout.js
+++ b/jstests/aggregation/bugs/cursor_timeout.js
@@ -94,5 +94,34 @@
         },
     ]);
 
+    // Test that an aggregation cursor with nested $lookup stages is killed when the timeout is
+    // reached.
+    assertCursorTimesOut('source', [
+        {
+          $lookup: {
+              from: 'dest',
+              let : {local1: "$local"},
+              pipeline: [
+                  {$match: {$expr: {$eq: ["$foreign", "$$local1"]}}},
+                  {
+                    $lookup: {
+                        from: 'source',
+                        let : {foreign1: "$foreign"},
+                        pipeline: [{$match: {$expr: {$eq: ["$local", "$$foreign1"]}}}],
+                        as: 'matches2'
+                    }
+                  },
+                  {
+                    $unwind: "$matches2",
+                  },
+              ],
+              as: 'matches1',
+          }
+        },
+        {
+          $unwind: "$matches1",
+        },
+    ]);
+
     MongoRunner.stopMongod(conn);
 })();
diff --git a/jstests/aggregation/bugs/lookup_unwind_getmore.js b/jstests/aggregation/bugs/lookup_unwind_getmore.js
index 1e48f93302a..9eaece9bfe8 100644
--- a/jstests/aggregation/bugs/lookup_unwind_getmore.js
+++ b/jstests/aggregation/bugs/lookup_unwind_getmore.js
@@ -13,25 +13,43 @@
 
     const testDB = conn.getDB('test');
 
-    // We use a batch size of 2 to ensure that the mongo shell issues a getMore when unwinding the
-    // results from the 'dest' collection for the same document in the 'source' collection under a
-    // different OperationContext.
-    const batchSize = 2;
-    const numMatches = 5;
-
-    assert.writeOK(testDB.source.insert({local: 1}));
-    for (let i = 0; i < numMatches; ++i) {
-        assert.writeOK(testDB.dest.insert({foreign: 1}));
+    /**
+     * Executes an aggregrate with 'options.pipeline' and confirms that 'options.numResults' were
+     * returned.
+     */
+    function runTest(options) {
+        // The batchSize must be smaller than the number of documents returned by the $lookup. This
+        // ensures that the mongo shell will issue a getMore when unwinding the $lookup results for
+        // the same document in the 'source' collection, under a different OperationContext.
+        const batchSize = 2;
+
+        testDB.source.drop();
+        assert.writeOK(testDB.source.insert({x: 1}));
+
+        testDB.dest.drop();
+        for (let i = 0; i < 5; ++i) {
+            assert.writeOK(testDB.dest.insert({x: 1}));
+        }
+
+        const res = assert.commandWorked(testDB.runCommand({
+            aggregate: 'source',
+            pipeline: options.pipeline,
+            cursor: {
+                batchSize: batchSize,
+            },
+        }));
+
+        const cursor = new DBCommandCursor(testDB, res, batchSize);
+        assert.eq(options.numResults, cursor.itcount());
     }
 
-    const res = assert.commandWorked(testDB.runCommand({
-        aggregate: 'source',
+    runTest({
         pipeline: [
             {
               $lookup: {
                   from: 'dest',
-                  localField: 'local',
-                  foreignField: 'foreign',
+                  localField: 'x',
+                  foreignField: 'x',
                   as: 'matches',
               }
             },
@@ -41,13 +59,42 @@
               },
             },
         ],
-        cursor: {
-            batchSize: batchSize,
-        },
-    }));
+        numResults: 5
+    });
 
-    const cursor = new DBCommandCursor(testDB, res, batchSize);
-    assert.eq(numMatches, cursor.itcount());
+    runTest({
+        pipeline: [
+            {
+              $lookup: {
+                  from: 'dest',
+                  let : {x1: "$x"},
+                  pipeline: [
+                      {$match: {$expr: {$eq: ["$$x1", "$x"]}}},
+                      {
+                        $lookup: {
+                            from: "dest",
+                            as: "matches2",
+                            let : {x2: "$x"},
+                            pipeline: [{$match: {$expr: {$eq: ["$$x2", "$x"]}}}]
+                        }
+                      },
+                      {
+                        $unwind: {
+                            path: '$matches2',
+                        },
+                      },
+                  ],
+                  as: 'matches1',
+              }
+            },
+            {
+              $unwind: {
+                  path: '$matches1',
+              },
+            },
+        ],
+        numResults: 25
+    });
 
     MongoRunner.stopMongod(conn);
 })();
diff --git a/jstests/aggregation/bugs/lookup_unwind_killcursor.js b/jstests/aggregation/bugs/lookup_unwind_killcursor.js
index 428284661d5..503bbf70e8d 100644
--- a/jstests/aggregation/bugs/lookup_unwind_killcursor.js
+++ b/jstests/aggregation/bugs/lookup_unwind_killcursor.js
@@ -13,43 +13,80 @@
 
     const testDB = conn.getDB('test');
 
-    // We use a batch size of 2 to ensure that the mongo shell does not exhaust the cursor on its
-    // first batch.
-    const batchSize = 2;
-    const numMatches = 5;
-
-    assert.writeOK(testDB.source.insert({local: 1}));
-    for (let i = 0; i < numMatches; ++i) {
-        assert.writeOK(testDB.dest.insert({foreign: 1}));
-    }
+    function runTest(pipeline) {
+        // We use a batch size of 2 to ensure that the mongo shell does not exhaust the cursor on
+        // its first batch.
+        const batchSize = 2;
 
-    const res = assert.commandWorked(testDB.runCommand({
-        aggregate: 'source',
-        pipeline: [
-            {
-              $lookup: {
-                  from: 'dest',
-                  localField: 'local',
-                  foreignField: 'foreign',
-                  as: 'matches',
-              }
-            },
-            {
-              $unwind: {
-                  path: '$matches',
-              },
+        testDB.source.drop();
+        assert.writeOK(testDB.source.insert({x: 1}));
+
+        testDB.dest.drop();
+        for (let i = 0; i < 5; ++i) {
+            assert.writeOK(testDB.dest.insert({x: 1}));
+        }
+
+        const res = assert.commandWorked(testDB.runCommand({
+            aggregate: 'source',
+            pipeline: pipeline,
+            cursor: {
+                batchSize: batchSize,
             },
-        ],
-        cursor: {
-            batchSize: batchSize,
-        },
-    }));
+        }));
+
+        const cursor = new DBCommandCursor(testDB, res, batchSize);
+        cursor.close();  // Closing the cursor will issue the "killCursors" command.
 
-    const cursor = new DBCommandCursor(testDB, res, batchSize);
-    cursor.close();  // Closing the cursor will issue the "killCursors" command.
+        const serverStatus = assert.commandWorked(testDB.adminCommand({serverStatus: 1}));
+        assert.eq(0, serverStatus.metrics.cursor.open.total, tojson(serverStatus.metrics.cursor));
+    }
+
+    runTest([
+        {
+          $lookup: {
+              from: 'dest',
+              localField: 'x',
+              foreignField: 'x',
+              as: 'matches',
+          }
+        },
+        {
+          $unwind: {
+              path: '$matches',
+          },
+        },
+    ]);
 
-    const serverStatus = assert.commandWorked(testDB.adminCommand({serverStatus: 1}));
-    assert.eq(0, serverStatus.metrics.cursor.open.total, tojson(serverStatus.metrics.cursor));
+    runTest([
+        {
+          $lookup: {
+              from: 'dest',
+              let : {x1: "$x"},
+              pipeline: [
+                  {$match: {$expr: {$eq: ["$$x1", "$x"]}}},
+                  {
+                    $lookup: {
+                        from: "dest",
+                        as: "matches2",
+                        let : {x2: "$x"},
+                        pipeline: [{$match: {$expr: {$eq: ["$$x2", "$x"]}}}]
+                    }
+                  },
+                  {
+                    $unwind: {
+                        path: '$matches2',
+                    },
+                  },
+              ],
+              as: 'matches1',
+          }
+        },
+        {
+          $unwind: {
+              path: '$matches1',
+          },
+        },
+    ]);
 
     MongoRunner.stopMongod(conn);
 })();
diff --git a/jstests/aggregation/sources/facet/inner_lookup.js b/jstests/aggregation/sources/facet/inner_lookup.js
index c53bfa2b48e..39f1b53f88f 100644
--- a/jstests/aggregation/sources/facet/inner_lookup.js
+++ b/jstests/aggregation/sources/facet/inner_lookup.js
@@ -21,16 +21,43 @@
     assert.writeOK(foreign.insert({_id: 1, foreignKey: 1}));
     assert.writeOK(foreign.insert({_id: 2, foreignKey: 2}));
 
-    const lookupStage = {
-        $lookup:
-            {from: foreign.getName(), localField: "_id", foreignField: "foreignKey", as: "joined"}
-    };
-    const lookupResults = local.aggregate([lookupStage]).toArray();
-    const facetedLookupResults = local.aggregate([{$facet: {nested: [lookupStage]}}]).toArray();
-    assert.eq(facetedLookupResults, [{nested: lookupResults}]);
-
-    const lookupResultsUnwound = local.aggregate([lookupStage, {$unwind: "$joined"}]).toArray();
-    const facetedLookupResultsUnwound =
-        local.aggregate([{$facet: {nested: [lookupStage, {$unwind: "$joined"}]}}]).toArray();
-    assert.eq(facetedLookupResultsUnwound, [{nested: lookupResultsUnwound}]);
+    function runTest(lookupStage) {
+        const lookupResults = local.aggregate([lookupStage]).toArray();
+        const facetedLookupResults = local.aggregate([{$facet: {nested: [lookupStage]}}]).toArray();
+        assert.eq(facetedLookupResults, [{nested: lookupResults}]);
+
+        const lookupResultsUnwound = local.aggregate([lookupStage, {$unwind: "$joined"}]).toArray();
+        const facetedLookupResultsUnwound =
+            local.aggregate([{$facet: {nested: [lookupStage, {$unwind: "$joined"}]}}]).toArray();
+        assert.eq(facetedLookupResultsUnwound, [{nested: lookupResultsUnwound}]);
+    }
+
+    runTest({
+        $lookup: {
+            from: foreign.getName(),
+            localField: "_id",
+            foreignField: "foreignKey",
+            as: "joined"
+        }
+    });
+
+    runTest({
+        $lookup: {
+            from: foreign.getName(),
+            let : {id1: "$_id"},
+            pipeline: [
+                {$match: {$expr: {$eq: ["$$id1", "$foreignKey"]}}},
+                {
+                  $lookup: {
+                      from: foreign.getName(),
+                      let : {id2: "$_id"},
+                      pipeline: [{$match: {$expr: {$eq: ["$$id2", "$foreignKey"]}}}],
+                      as: "joined2"
+                  }
+                },
+                {$unwind: "$joined2"}
+            ],
+            as: "joined"
+        }
+    });
 }());
diff --git a/jstests/aggregation/sources/lookup/collation_lookup.js b/jstests/aggregation/sources/lookup/collation_lookup.js
index 89a0b21c038..80d173138a6 100644
--- a/jstests/aggregation/sources/lookup/collation_lookup.js
+++ b/jstests/aggregation/sources/lookup/collation_lookup.js
@@ -46,6 +46,44 @@
         arrayEq(expected, res[0].matched),
         "Expected " + tojson(expected) + " to equal " + tojson(res[0].matched) + " up to ordering");
 
+    res = withDefaultCollationColl
+              .aggregate([{
+                  $lookup: {
+                      from: withoutDefaultCollationColl.getName(),
+                      let : {str1: "$str"},
+                      pipeline: [
+                          {$match: {$expr: {$eq: ["$str", "$$str1"]}}},
+                          {
+                            $lookup: {
+                                from: withoutDefaultCollationColl.getName(),
+                                let : {str2: "$str"},
+                                pipeline: [{$match: {$expr: {$eq: ["$str", "$$str1"]}}}],
+                                as: "matched2"
+                            }
+                          }
+                      ],
+                      as: "matched1",
+                  },
+              }])
+              .toArray();
+    assert.eq(1, res.length, tojson(res));
+
+    expected = [
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched2": [{"_id": "lowercase", "str": "abc"}, {"_id": "uppercase", "str": "ABC"}]
+        },
+        {
+          "_id": "uppercase",
+          "str": "ABC",
+          "matched2": [{"_id": "lowercase", "str": "abc"}, {"_id": "uppercase", "str": "ABC"}]
+        }
+    ];
+    assert(arrayEq(expected, res[0].matched1),
+           "Expected " + tojson(expected) + " to equal " + tojson(res[0].matched1) +
+               " up to ordering. " + tojson(res));
+
     // Test that the $lookup stage respects the inherited collation when it optimizes with an
     // $unwind stage.
     res = withDefaultCollationColl
@@ -70,6 +108,61 @@
     assert(arrayEq(expected, res),
            "Expected " + tojson(expected) + " to equal " + tojson(res) + " up to ordering");
 
+    res = withDefaultCollationColl
+              .aggregate([
+                  {
+                    $lookup: {
+                        from: withoutDefaultCollationColl.getName(),
+                        let : {str1: "$str"},
+                        pipeline: [
+                            {$match: {$expr: {$eq: ["$str", "$$str1"]}}},
+                            {
+                              $lookup: {
+                                  from: withoutDefaultCollationColl.getName(),
+                                  let : {str2: "$str"},
+                                  pipeline: [{$match: {$expr: {$eq: ["$str", "$$str1"]}}}],
+                                  as: "matched2"
+                              }
+                            },
+                            {$unwind: "$matched2"},
+                        ],
+                        as: "matched1",
+                    },
+                  },
+                  {$unwind: "$matched1"},
+              ])
+              .toArray();
+    assert.eq(4, res.length, tojson(res));
+
+    expected = [
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "lowercase", "str": "abc", "matched2": {"_id": "lowercase", "str": "abc"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "lowercase", "str": "abc", "matched2": {"_id": "uppercase", "str": "ABC"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "uppercase", "str": "ABC", "matched2": {"_id": "lowercase", "str": "abc"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "uppercase", "str": "ABC", "matched2": {"_id": "uppercase", "str": "ABC"}}
+        }
+    ];
+    assert(arrayEq(expected, res),
+           "Expected " + tojson(expected) + " to equal " + tojson(res) + " up to ordering");
+
     // Test that the $lookup stage respects an explicit collation on the aggregation operation.
     res = withoutDefaultCollationColl
               .aggregate(
@@ -93,6 +186,49 @@
         arrayEq(expected, res[0].matched),
         "Expected " + tojson(expected) + " to equal " + tojson(res[0].matched) + " up to ordering");
 
+    res = withoutDefaultCollationColl
+              .aggregate(
+                  [
+                    {$match: {_id: "lowercase"}},
+                    {
+                      $lookup: {
+                          from: withoutDefaultCollationColl.getName(),
+                          let : {str1: "$str"},
+                          pipeline: [
+                              {$match: {$expr: {$eq: ["$str", "$$str1"]}}},
+                              {
+                                $lookup: {
+                                    from: withoutDefaultCollationColl.getName(),
+                                    let : {str2: "$str"},
+                                    pipeline: [{$match: {$expr: {$eq: ["$str", "$$str1"]}}}],
+                                    as: "matched2"
+                                }
+                              }
+                          ],
+                          as: "matched1",
+                      },
+                    }
+                  ],
+                  caseInsensitive)
+              .toArray();
+    assert.eq(1, res.length, tojson(res));
+
+    expected = [
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched2": [{"_id": "lowercase", "str": "abc"}, {"_id": "uppercase", "str": "ABC"}]
+        },
+        {
+          "_id": "uppercase",
+          "str": "ABC",
+          "matched2": [{"_id": "lowercase", "str": "abc"}, {"_id": "uppercase", "str": "ABC"}]
+        }
+    ];
+    assert(arrayEq(expected, res[0].matched1),
+           "Expected " + tojson(expected) + " to equal " + tojson(res[0].matched1) +
+               " up to ordering");
+
     // Test that the $lookup stage respects an explicit collation on the aggregation operation when
     // it optimizes with an $unwind stage.
     res = withoutDefaultCollationColl
@@ -120,6 +256,64 @@
     assert(arrayEq(expected, res),
            "Expected " + tojson(expected) + " to equal " + tojson(res) + " up to ordering");
 
+    res = withoutDefaultCollationColl
+              .aggregate(
+                  [
+                    {$match: {_id: "lowercase"}},
+                    {
+                      $lookup: {
+                          from: withoutDefaultCollationColl.getName(),
+                          let : {str1: "$str"},
+                          pipeline: [
+                              {$match: {$expr: {$eq: ["$str", "$$str1"]}}},
+                              {
+                                $lookup: {
+                                    from: withoutDefaultCollationColl.getName(),
+                                    let : {str2: "$str"},
+                                    pipeline: [{$match: {$expr: {$eq: ["$str", "$$str1"]}}}],
+                                    as: "matched2"
+                                }
+                              },
+                              {$unwind: "$matched2"},
+                          ],
+                          as: "matched1",
+                      },
+                    },
+                    {$unwind: "$matched1"},
+                  ],
+                  caseInsensitive)
+              .toArray();
+    assert.eq(4, res.length, tojson(res));
+
+    expected = [
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "lowercase", "str": "abc", "matched2": {"_id": "lowercase", "str": "abc"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "lowercase", "str": "abc", "matched2": {"_id": "uppercase", "str": "ABC"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "uppercase", "str": "ABC", "matched2": {"_id": "lowercase", "str": "abc"}}
+        },
+        {
+          "_id": "lowercase",
+          "str": "abc",
+          "matched1":
+              {"_id": "uppercase", "str": "ABC", "matched2": {"_id": "uppercase", "str": "ABC"}}
+        }
+    ];
+    assert(arrayEq(expected, res),
+           "Expected " + tojson(expected) + " to equal " + tojson(res) + " up to ordering");
+
     // Test that the $lookup stage uses the "simple" collation if a collation isn't set on the
     // collection or the aggregation operation.
     res = withoutDefaultCollationColl
@@ -136,4 +330,39 @@
               ])
               .toArray();
     assert.eq([{_id: "lowercase", str: "abc", matched: [{_id: "lowercase", str: "abc"}]}], res);
+
+    res = withoutDefaultCollationColl
+              .aggregate([
+                  {$match: {_id: "lowercase"}},
+                  {
+                    $lookup: {
+                        from: withoutDefaultCollationColl.getName(),
+                        let : {str1: "$str"},
+                        pipeline: [
+                            {$match: {$expr: {$eq: ["$str", "$$str1"]}}},
+                            {
+                              $lookup: {
+                                  from: withoutDefaultCollationColl.getName(),
+                                  let : {str2: "$str"},
+                                  pipeline: [{$match: {$expr: {$eq: ["$str", "$$str1"]}}}],
+                                  as: "matched2"
+                              }
+                            },
+                            {$unwind: "$matched2"},
+                        ],
+                        as: "matched1",
+                    },
+                  },
+              ])
+              .toArray();
+    assert.eq([{
+                 "_id": "lowercase",
+                 "str": "abc",
+                 "matched1": [{
+                     "_id": "lowercase",
+                     "str": "abc",
+                     "matched2": {"_id": "lowercase", "str": "abc"}
+                 }]
+              }],
+              res);
 })();