SERVER-66548 Add support for dependency tracking to $redact

(cherry picked from commit cc4968771221658d66ff42ed2fb861656c5683ca)

3 files changed, 66 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/lookup/lookup_non_correlated_prefix.js b/jstests/aggregation/sources/lookup/lookup_non_correlated_prefix.js
index b4feed0ab31..d377d1016cc 100644
--- a/jstests/aggregation/sources/lookup/lookup_non_correlated_prefix.js
+++ b/jstests/aggregation/sources/lookup/lookup_non_correlated_prefix.js
@@ -8,6 +8,8 @@
 (function() {
 "use strict";
 
+load("jstests/aggregation/extras/utils.js");  // for arrayEq
+
 const testColl = db.lookup_non_correlated_prefix;
 testColl.drop();
 const joinColl = db.lookup_non_correlated_prefix_join;
@@ -78,6 +80,39 @@ cursor.toArray().forEach(user => {
     assert.eq(user['_id'], joinedDocs[0].owner);
 });
 
+// Test for a non-correlated prefix followed by a $facet pipeline that contains a correlated
+// variable reference.
+cursor = testColl.aggregate([
+    {
+        $lookup: {
+            as: 'items_check',
+            from: joinColl.getName(),
+            let : {id: '$_id'},
+            pipeline: [
+                {$match: {owner: "user_1"}},
+                {
+                    $facet: {
+                        all: [{
+                            $redact: {
+                                $cond:
+                                    {if: {$eq: ["$$id", "user_1"]}, then: "$$KEEP", else: "$$PRUNE"}
+                            }
+                        }],
+                    },
+                },
+            ],
+        },
+    },
+]);
+res = cursor.toArray();
+assert(
+    arrayEq(res,
+            [
+                {"_id": "user_1", "items_check": [{"all": [{"_id": "item_1", "owner": "user_1"}]}]},
+                {"_id": "user_2", "items_check": [{"all": []}]}
+            ]),
+    res);
+
 // SERVER-57000: Test handling of lack of correlation (addFields with empty set of columns)
 assert.doesNotThrow(() => testColl.aggregate([
     {
diff --git a/src/mongo/db/pipeline/document_source_redact.h b/src/mongo/db/pipeline/document_source_redact.h
index a45e2af8c9f..cf75fa00a51 100644
--- a/src/mongo/db/pipeline/document_source_redact.h
+++ b/src/mongo/db/pipeline/document_source_redact.h
@@ -72,6 +72,16 @@ public:
         return _expression;
     }
 
+    DepsTracker::State getDependencies(DepsTracker* deps) const final {
+        // Add the dependencies of the expression but all we really care about is variable
+        // references for correlation analysis. The field references may get populated but we'll
+        // still require the full document since the $redact may descend arbitrary levels of nested
+        // documents that is only known at runtime.
+        _expression->addDependencies(deps);
+        deps->needWholeDocument = true;
+        return DepsTracker::State::SEE_NEXT;
+    }
+
 private:
     DocumentSourceRedact(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                          const boost::intrusive_ptr<Expression>& previsit);
diff --git a/src/mongo/db/pipeline/document_source_redact_test.cpp b/src/mongo/db/pipeline/document_source_redact_test.cpp
index 1c4ddf60f92..92d8bc08766 100644
--- a/src/mongo/db/pipeline/document_source_redact_test.cpp
+++ b/src/mongo/db/pipeline/document_source_redact_test.cpp
@@ -80,5 +80,26 @@ TEST_F(DocumentSourceRedactTest, ShouldPropagatePauses) {
     ASSERT_TRUE(redact->getNext().isEOF());
     ASSERT_TRUE(redact->getNext().isEOF());
 }
+
+TEST_F(DocumentSourceRedactTest, ReportsVariableDependencies) {
+    auto varId = getExpCtx()->variablesParseState.defineVariable("var");
+    auto redactSpec = fromjson(R"({
+        "$redact" : {
+            "$cond" : {
+                "if" : "$$var",
+                "then" : "$$PRUNE",
+                "else" : "$$DESCEND"
+            }
+        }
+    })");
+    auto redact = DocumentSourceRedact::createFromBson(redactSpec.firstElement(), getExpCtx());
+
+    DepsTracker deps;
+    ASSERT_EQ(redact->getDependencies(&deps), DepsTracker::State::SEE_NEXT);
+    ASSERT_EQ(deps.needWholeDocument, true);
+    ASSERT_TRUE(deps.fields.empty());
+    ASSERT_EQ(deps.vars.count(varId), 1);
+}
+
 }  // namespace
 }  // namespace mongo