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

3 files changed, 111 insertions, 5 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index 8c5b74707cd..906237ea4ef 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -74,6 +74,8 @@ last-continuous:
       test_file: jstests/sharding/range_deleter_interacts_correctly_with_refine_shard_key.js
     - ticket: SERVER-61977
       test_file: jstests/replsets/rollback_during_step_up.js
+    - ticket: SERVER-63141
+      test_file: jstests/aggregation/lookup_let_optimization.js
 
   # Tests that should only be excluded from particular suites should be listed under that suite.
   suites:
@@ -304,6 +306,8 @@ last-lts:
       test_file: jstests/sharding/range_deleter_interacts_correctly_with_refine_shard_key.js
     - ticket: SERVER-61977
       test_file: jstests/replsets/rollback_during_step_up.js
+    - ticket: SERVER-63141
+      test_file: jstests/aggregation/lookup_let_optimization.js
 
   # Tests that should only be excluded from particular suites should be listed under that suite.
   suites:
diff --git a/jstests/aggregation/lookup_let_optimization.js b/jstests/aggregation/lookup_let_optimization.js
new file mode 100644
index 00000000000..c9d43e80f8f
--- /dev/null
+++ b/jstests/aggregation/lookup_let_optimization.js
@@ -0,0 +1,101 @@
+/**
+ * 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
+ * ]
+ */
+
+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 a91be27fdd7..33c32f096b3 100644
--- a/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp
+++ b/src/mongo/db/pipeline/document_source_sequential_document_cache.cpp
@@ -118,13 +118,14 @@ Pipeline::SourceContainer::iterator DocumentSourceSequentialDocumentCache::doOpt
     DepsTracker deps(DepsTracker::kNoMetadata);
 
     // Iterate through the pipeline stages until we find one which cannot be cached.
-    // A stage cannot be cached if it either: 1. depends on a variable defined in this scope, or
-    // 2. generates random numbers.
+    // A stage cannot be cached if it either:
+    //  1. does not support dependency tracking, and may thus require the full object and metadata.
+    //  2. depends on a variable defined in this scope, or
+    //  3. generates random numbers.
     DocumentSource* lastPtr = nullptr;
     for (; prefixSplit != container->end(); ++prefixSplit) {
-        (*prefixSplit)->getDependencies(&deps);
-
-        if (deps.hasVariableReferenceTo(varIDs) || deps.needRandomGenerator) {
+        if (((*prefixSplit)->getDependencies(&deps) == DepsTracker::State::NOT_SUPPORTED) ||
+            deps.hasVariableReferenceTo(varIDs) || deps.needRandomGenerator) {
             break;
         }