SERVER-46998 Support 'let' variables used within $merge custom pipeline update

4 files changed, 103 insertions, 21 deletions

diff --git a/jstests/noPassthroughWithMongod/command_let_variables.js b/jstests/noPassthroughWithMongod/command_let_variables.js
index aa04bedbf62..1b80977f0dd 100644
--- a/jstests/noPassthroughWithMongod/command_let_variables.js
+++ b/jstests/noPassthroughWithMongod/command_let_variables.js
@@ -6,7 +6,8 @@
 (function() {
 "use strict";
 
-const coll = db.update_let_variables;
+const coll = db.command_let_variables;
+const targetColl = db.command_let_variables_target;
 coll.drop();
 
 assert.commandWorked(coll.insert([
@@ -110,6 +111,91 @@ assert.commandFailedWithCode(db.runCommand({
 }),
                              ErrorCodes.TypeMismatch);
 
+// Function to prepare target collection of $merge stage for testing.
+function prepMergeTargetColl() {
+    targetColl.drop();
+
+    assert.commandWorked(db.runCommand({
+        aggregate: coll.getName(),
+        pipeline: [
+            {$match: {$expr: {$eq: ["$Species", "Song Thrush (Turdus philomelos)"]}}},
+            {$out: targetColl.getName()}
+        ],
+        cursor: {}
+    }));
+}
+
+// Test that $merge stage can use 'let' variables within its own stage's pipeline.
+prepMergeTargetColl();
+assert.commandWorked(db.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $merge: {
+            into: targetColl.getName(),
+            let : {variable: "INNER"},
+            whenMatched: [{$addFields: {"var": "$$variable"}}]
+        }
+    }],
+    cursor: {}
+}));
+assert.eq(targetColl.aggregate({$match: {$expr: {$eq: ["$var", "INNER"]}}}).toArray().length, 1);
+
+// Test that $merge stage can access command-level 'let' variables.
+prepMergeTargetColl();
+assert.commandWorked(db.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [
+        {$merge: {into: targetColl.getName(), whenMatched: [{$addFields: {"var": "$$variable"}}]}}
+    ],
+    cursor: {},
+    let : {variable: "OUTER"}
+}));
+assert.eq(targetColl.aggregate({$match: {$expr: {$eq: ["$var", "OUTER"]}}}).toArray().length, 1);
+
+// Test that $merge stage can use stage-level and command-level 'let' variables in same command.
+prepMergeTargetColl();
+assert.commandWorked(db.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $merge: {
+            into: targetColl.getName(),
+            let : {stage: "INNER"},
+            whenMatched: [{$addFields: {"innerVar": "$$stage", "outerVar": "$$command"}}]
+        }
+    }],
+    cursor: {},
+    let : {command: "OUTER"}
+}));
+assert.eq(
+    targetColl
+        .aggregate({
+            $match: {
+                $and:
+                    [{$expr: {$eq: ["$innerVar", "INNER"]}},
+                     {$expr: {$eq: ["$outerVar", "OUTER"]}}]
+            }
+        })
+        .toArray()
+        .length,
+    1);
+
+// Test that $merge stage follows variable scoping rules with stage-level and command-level 'let'
+// variables.
+prepMergeTargetColl();
+assert.commandWorked(db.runCommand({
+    aggregate: coll.getName(),
+    pipeline: [{
+        $merge: {
+            into: targetColl.getName(),
+            let : {variable: "INNER"},
+            whenMatched: [{$addFields: {"var": "$$variable"}}]
+        }
+    }],
+    cursor: {},
+    let : {variable: "OUTER"}
+}));
+assert.eq(targetColl.aggregate({$match: {$expr: {$eq: ["$var", "INNER"]}}}).toArray().length, 1);
+
 // findAndModify
 assert.commandWorked(coll.insert({Species: "spy_bird"}));
 let result = db.runCommand({
diff --git a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
index 8f4c9fbb392..0870cbd0590 100644
--- a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
@@ -613,7 +613,10 @@ Update CommonMongodProcessInterface::buildUpdateOp(
         return wcb;
     }());
     updateOp.setRuntimeConstants(expCtx->getRuntimeConstants());
-    updateOp.setLet(expCtx->variables.serializeLetParameters(expCtx->variablesParseState));
+    if (auto letParams = expCtx->variablesParseState.serializeUserVariables(expCtx->variables);
+        !letParams.isEmpty()) {
+        updateOp.setLet(letParams);
+    }
     return updateOp;
 }
 
diff --git a/src/mongo/db/pipeline/variables.cpp b/src/mongo/db/pipeline/variables.cpp
index b64ec45db00..4ae50362250 100644
--- a/src/mongo/db/pipeline/variables.cpp
+++ b/src/mongo/db/pipeline/variables.cpp
@@ -142,7 +142,7 @@ void Variables::setValue(Id id, const Value& value, bool isConstant) {
     // If a value has already been set for 'id', and that value was marked as constant, then it
     // is illegal to modify.
     invariant(!hasConstantValue(id));
-    _values[id] = {value, isConstant};
+    _letParametersMap[id] = {value, isConstant};
 }
 
 void Variables::setValue(Variables::Id id, const Value& value) {
@@ -158,8 +158,8 @@ void Variables::setConstantValue(Variables::Id id, const Value& value) {
 Value Variables::getUserDefinedValue(Variables::Id id) const {
     invariant(isUserDefinedVariable(id));
 
-    auto it = _values.find(id);
-    uassert(40434, str::stream() << "Undefined variable id: " << id, it != _values.end());
+    auto it = _letParametersMap.find(id);
+    uassert(40434, str::stream() << "Undefined variable id: " << id, it != _letParametersMap.end());
     return it->second.value;
 }
 
@@ -233,12 +233,6 @@ void Variables::setDefaultRuntimeConstants(OperationContext* opCtx) {
     setRuntimeConstants(Variables::generateRuntimeConstants(opCtx));
 }
 
-BSONObj Variables::serializeLetParameters(const VariablesParseState& vps) const {
-    auto bob = BSONObjBuilder{};
-    for (auto&& [id, value] : _letParametersMap)
-        bob << kIdToBuiltinVarName.at(id) << value;
-    return bob.appendElements(vps.serialize(*this)).obj();
-}
 void Variables::seedVariablesWithLetParameters(ExpressionContext* const expCtx,
                                                const BSONObj letParams) {
     for (auto&& elem : letParams) {
@@ -332,7 +326,7 @@ std::set<Variables::Id> VariablesParseState::getDefinedVariableIDs() const {
     return ids;
 }
 
-BSONObj VariablesParseState::serialize(const Variables& vars) const {
+BSONObj VariablesParseState::serializeUserVariables(const Variables& vars) const {
     auto bob = BSONObjBuilder{};
     for (auto&& [var_name, id] : _variables)
         if (vars.hasValue(id))
diff --git a/src/mongo/db/pipeline/variables.h b/src/mongo/db/pipeline/variables.h
index b868f51f4e0..6cf5256a640 100644
--- a/src/mongo/db/pipeline/variables.h
+++ b/src/mongo/db/pipeline/variables.h
@@ -128,7 +128,8 @@ public:
      * Returns whether a constant value for 'id' has been defined using setConstantValue().
      */
     bool hasConstantValue(Variables::Id id) const {
-        if (auto it = _values.find(id); it != _values.end() && it->second.isConstant) {
+        if (auto it = _letParametersMap.find(id);
+            it != _letParametersMap.end() && it->second.isConstant) {
             return true;
         }
         return false;
@@ -162,11 +163,6 @@ public:
     void setDefaultRuntimeConstants(OperationContext* opCtx);
 
     /**
-     * Return an object which represents the variables which are considered let parameters.
-     */
-    BSONObj serializeLetParameters(const VariablesParseState& vps) const;
-
-    /**
      * Seed let parameters with the given BSONObj.
      */
     void seedVariablesWithLetParameters(ExpressionContext* const expCtx,
@@ -218,9 +214,8 @@ private:
     }
 
     IdGenerator _idGenerator;
-    stdx::unordered_map<Id, ValueAndState> _values;
+    stdx::unordered_map<Id, ValueAndState> _letParametersMap;
     stdx::unordered_map<Id, Value> _runtimeConstantsMap;
-    stdx::unordered_map<Id, Value> _letParametersMap;
 
     // Populated after construction. Should not be set more than once.
     boost::optional<RuntimeConstants> _runtimeConstants;
@@ -267,7 +262,11 @@ public:
      */
     std::set<Variables::Id> getDefinedVariableIDs() const;
 
-    BSONObj serialize(const Variables& vars) const;
+    /**
+     * Serializes the IDs and associated values of user-defined variables that are currently in
+     * scope.
+     */
+    BSONObj serializeUserVariables(const Variables& vars) const;
 
     /**
      * Return a copy of this VariablesParseState. Will replace the copy's '_idGenerator' pointer