SERVER-63141 Ensure stages without dependency tracking are not cached in $lookup

3 files changed, 106 insertions, 3 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 580840854e3..f5b8c3616ef 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -197,6 +197,8 @@ all:
     test_file: jstests/core/exhaust.js
   - ticket: SERVER-61977
     test_file: jstests/replsets/rollback_during_step_up.js
+  - ticket: SERVER-63141
+    test_file: jstests/aggregation/lookup_let_optimization.js
 
 
 suites:
diff --git a/jstests/aggregation/lookup_let_optimization.js b/jstests/aggregation/lookup_let_optimization.js
new file mode 100644
index 00000000000..f9aa3b5cf1c
--- /dev/null
+++ b/jstests/aggregation/lookup_let_optimization.js
@@ -0,0 +1,102 @@
+/**
+ * This test ensures that stages dependent on a let variable optimizing to a constant in a $lookup
+ * pipeline are evaluated correctly.
+ * @tags: [
+ *     requires_pipeline_optimization,
+ *     assumes_unsharded_collection,
+ *     assumes_against_mongod_not_mongos,
+ * ]
+ */
+
+load('jstests/aggregation/extras/utils.js');  // For assertArrayEq.
+load("jstests/libs/fixture_helpers.js");      // For FixtureHelpers.
+
+(function() {
+"use strict";
+
+const collName = "lookup_let_redact";
+const coll = db[collName];
+coll.drop();
+assert.commandWorked(coll.insert([
+    {_id: "true", test: true},
+    {_id: "false", test: false},
+]));
+
+const admin = db.getSiblingDB("admin");
+
+const setPipelineOptimizationMode = (mode) => {
+    FixtureHelpers.runCommandOnEachPrimary(
+        {db: admin, cmdObj: {configureFailPoint: 'disablePipelineOptimization', mode}});
+};
+
+const verifyAggregationForBothPipelineOptimizationModes = ({pipeline, expected}) => {
+    // Verify results when pipeline optimization is disabled.
+    setPipelineOptimizationMode('alwaysOn');
+    assertArrayEq({actual: coll.aggregate(pipeline).toArray(), expected});
+
+    // Verify results when pipeline optimization is enabled.
+    setPipelineOptimizationMode('off');
+    assertArrayEq({actual: coll.aggregate(pipeline).toArray(), expected});
+};
+
+// Get initial optimization mode.
+const pipelineOptParameter = assert.commandWorked(
+    db.adminCommand({getParameter: 1, "failpoint.disablePipelineOptimization": 1}));
+const oldMode =
+    pipelineOptParameter["failpoint.disablePipelineOptimization"].mode ? 'alwaysOn' : 'off';
+
+// Verify $redact.
+verifyAggregationForBothPipelineOptimizationModes({
+    pipeline: [
+        {$lookup: {
+            from: collName,
+            let: {iShouldPrune: "$test"},
+            pipeline: [
+                {$redact: {$cond: {if: "$$iShouldPrune", then: "$$PRUNE", else: "$$DESCEND"}}}
+            ],
+            as: "redacted"
+        }}
+    ],
+    expected: [
+        {_id: "true", test: true, redacted: []}, // Expect that documents were pruned.
+        {_id: "false", test: false, redacted: [ // Expect that $redact descended instead.
+            {_id: "true", test: true},
+            {_id: "false", test: false} 
+        ]}
+    ]
+});
+
+// Verify $unionWith.
+verifyAggregationForBothPipelineOptimizationModes({
+    pipeline: [
+        {$lookup: {
+            from: collName,
+            let: {iShouldMatch: "$test"},
+            pipeline: [
+                {$unionWith: {
+                    coll: collName,
+                    pipeline: [{$match: {$expr: {$eq: ["$$iShouldMatch", "$test"]}}}]
+                }}
+            ],
+            as: "united"
+        }}
+    ],
+    expected: [
+        // $unionWith should match and include only the document with {test: true} in 'united'.
+        {_id: "true", test: true, united: [
+            {_id: "true", test: true},
+            {_id: "true", test: true},
+            {_id: "false", test: false},
+        ]},
+        // $unionWith should match and include only the document with {test: false} in 'united'.
+        {_id: "false", test: false, united: [
+            {_id: "false", test: false},
+            {_id: "false", test: false},
+            {_id: "true", test: true},
+        ]}
+    ]
+});
+
+// Reset optimization mode.
+setPipelineOptimizationMode(oldMode);
+}());
diff --git a/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp b/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp
index a23563adca5..5b467f45cf4 100644
--- a/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp
+++ b/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp
@@ -119,9 +119,8 @@ Pipeline::SourceContainer::iterator DocumentSourceSequentialDocumentCache::doOpt
 
     // Iterate through the pipeline stages until we find one which references an external variable.
     for (; prefixSplit != container->end(); ++prefixSplit) {
-        (*prefixSplit)->getDependencies(&deps);
-
-        if (deps.hasVariableReferenceTo(varIDs)) {
+        if (((*prefixSplit)->getDependencies(&deps) == DepsTracker::State::NOT_SUPPORTED) ||
+            deps.hasVariableReferenceTo(varIDs)) {
             break;
         }
     }