SERVER-46708 Support let variables in update command

4 files changed, 42 insertions, 10 deletions

diff --git a/src/mongo/db/ops/parsed_update.cpp b/src/mongo/db/ops/parsed_update.cpp
index cdcc2428029..18285ed37d5 100644
--- a/src/mongo/db/ops/parsed_update.cpp
+++ b/src/mongo/db/ops/parsed_update.cpp
@@ -45,8 +45,11 @@ ParsedUpdate::ParsedUpdate(OperationContext* opCtx,
                            const ExtensionsCallback& extensionsCallback)
     : _opCtx(opCtx),
       _request(request),
-      _expCtx(make_intrusive<ExpressionContext>(
-          opCtx, nullptr, _request->getNamespaceString(), _request->getRuntimeConstants())),
+      _expCtx(make_intrusive<ExpressionContext>(opCtx,
+                                                nullptr,
+                                                _request->getNamespaceString(),
+                                                _request->getRuntimeConstants(),
+                                                _request->getLetParameters())),
       _driver(_expCtx),
       _canonicalQuery(),
       _extensionsCallback(extensionsCallback) {}
@@ -146,10 +149,14 @@ Status ParsedUpdate::parseQueryToCQ() {
         allowedMatcherFeatures &= ~MatchExpressionParser::AllowedFeatures::kExpr;
     }
 
-    // If the update request has runtime constants attached to it, pass them to the QueryRequest.
+    // If the update request has runtime constants or let parameters attached to it, pass them to
+    // the QueryRequest.
     if (auto& runtimeConstants = _request->getRuntimeConstants()) {
         qr->setRuntimeConstants(*runtimeConstants);
     }
+    if (auto& letParams = _request->getLetParameters()) {
+        qr->setLetParameters(*letParams);
+    }
 
     auto statusWithCQ = CanonicalQuery::canonicalize(
         _opCtx, std::move(qr), _expCtx, _extensionsCallback, allowedMatcherFeatures);
diff --git a/src/mongo/db/ops/update_request.h b/src/mongo/db/ops/update_request.h
index 5570f544c3e..6404b5dcd70 100644
--- a/src/mongo/db/ops/update_request.h
+++ b/src/mongo/db/ops/update_request.h
@@ -129,6 +129,14 @@ public:
         return _runtimeConstants;
     }
 
+    void setLetParameters(const boost::optional<BSONObj>& letParameters) {
+        _letParameters = letParameters;
+    }
+
+    const boost::optional<BSONObj>& getLetParameters() const {
+        return _letParameters;
+    }
+
     void setArrayFilters(const std::vector<BSONObj>& arrayFilters) {
         _updateOp.setArrayFilters(arrayFilters);
     }
@@ -287,6 +295,9 @@ private:
     // System-defined constant values which may be required by the query or update operation.
     boost::optional<RuntimeConstants> _runtimeConstants;
 
+    // User-defined constant values to be used with a pipeline-style update. These can be specified
+    // by the user for each individual element of the 'updates' array in the 'update' command.
+    boost::optional<BSONObj> _letParameters;
 
     // The statement id of this request.
     StmtId _stmtId = kUninitializedStmtId;
diff --git a/src/mongo/db/ops/write_ops.idl b/src/mongo/db/ops/write_ops.idl
index 464b25ae9d4..7279fc3bf1e 100644
--- a/src/mongo/db/ops/write_ops.idl
+++ b/src/mongo/db/ops/write_ops.idl
@@ -184,6 +184,11 @@ commands:
                 description: "An array of one or more update statements to perform."
                 type: array<UpdateOpEntry>
                 supports_doc_sequence: true
+            let:
+                description: "A set of user-specified constants used by pipeline-style update
+                              operations and $expr."
+                type: object
+                optional: true
             runtimeConstants:
                 description: "A collection of values that do not change once computed. These are
                               used by pipeline-style update operations."
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index 3ed033eb8e0..c405f64614e 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -708,11 +708,13 @@ static SingleWriteResult performSingleUpdateOp(OperationContext* opCtx,
 /**
  * Performs a single update, retrying failure due to DuplicateKeyError when eligible.
  */
-static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(OperationContext* opCtx,
-                                                              const NamespaceString& ns,
-                                                              StmtId stmtId,
-                                                              const write_ops::UpdateOpEntry& op,
-                                                              RuntimeConstants runtimeConstants) {
+static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(
+    OperationContext* opCtx,
+    const NamespaceString& ns,
+    StmtId stmtId,
+    const write_ops::UpdateOpEntry& op,
+    RuntimeConstants runtimeConstants,
+    const boost::optional<BSONObj>& letParams) {
     globalOpCounters.gotUpdate();
     ServerWriteConcernMetrics::get(opCtx)->recordWriteConcernForUpdate(opCtx->getWriteConcern());
     auto& curOp = *CurOp::get(opCtx);
@@ -732,6 +734,9 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(OperationContext*
     UpdateRequest request(op);
     request.setNamespaceString(ns);
     request.setRuntimeConstants(std::move(runtimeConstants));
+    if (letParams) {
+        request.setLetParameters(std::move(letParams));
+    }
     request.setStmtId(stmtId);
     request.setYieldPolicy(opCtx->inMultiDocumentTransaction() ? PlanExecutor::INTERRUPT_ONLY
                                                                : PlanExecutor::YIELD_AUTO);
@@ -817,8 +822,12 @@ WriteResult performUpdates(OperationContext* opCtx, const write_ops::Update& who
         ON_BLOCK_EXIT([&] { finishCurOp(opCtx, &curOp); });
         try {
             lastOpFixer.startingOp();
-            out.results.emplace_back(performSingleUpdateOpWithDupKeyRetry(
-                opCtx, wholeOp.getNamespace(), stmtId, singleOp, runtimeConstants));
+            out.results.emplace_back(performSingleUpdateOpWithDupKeyRetry(opCtx,
+                                                                          wholeOp.getNamespace(),
+                                                                          stmtId,
+                                                                          singleOp,
+                                                                          runtimeConstants,
+                                                                          wholeOp.getLet()));
             lastOpFixer.finishedOpSuccessfully();
         } catch (const DBException& ex) {
             const bool canContinue =