SERVER-46708 Support let variables in update command

11 files changed, 93 insertions, 10 deletions

diff --git a/src/mongo/db/commands/write_commands/write_commands.cpp b/src/mongo/db/commands/write_commands/write_commands.cpp
index 2e8d5f7e51d..44bf077f4d6 100644
--- a/src/mongo/db/commands/write_commands/write_commands.cpp
+++ b/src/mongo/db/commands/write_commands/write_commands.cpp
@@ -417,6 +417,7 @@ private:
             updateRequest.setNamespaceString(_batch.getNamespace());
             updateRequest.setRuntimeConstants(
                 _batch.getRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx)));
+            updateRequest.setLetParameters(_batch.getLet());
             updateRequest.setYieldPolicy(PlanExecutor::YIELD_AUTO);
             updateRequest.setExplain();
 
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 =
diff --git a/src/mongo/db/pipeline/expression_context.cpp b/src/mongo/db/pipeline/expression_context.cpp
index 22f071c078c..beedcb28752 100644
--- a/src/mongo/db/pipeline/expression_context.cpp
+++ b/src/mongo/db/pipeline/expression_context.cpp
@@ -140,6 +140,7 @@ ExpressionContext::ExpressionContext(OperationContext* opCtx,
                                      std::unique_ptr<CollatorInterface> collator,
                                      const NamespaceString& nss,
                                      const boost::optional<RuntimeConstants>& runtimeConstants,
+                                     const boost::optional<BSONObj>& letParameters,
                                      bool mayDbProfile)
     : ns(nss),
       opCtx(opCtx),
@@ -157,6 +158,16 @@ ExpressionContext::ExpressionContext(OperationContext* opCtx,
     }
 
     jsHeapLimitMB = internalQueryJavaScriptHeapSizeLimitMB.load();
+    if (letParameters) {
+        // TODO SERVER-47713: One possible fix is to change the interface of everything that needs
+        // an expression context intrusive_ptr to take a raw ptr.
+        auto intrusiveThis = boost::intrusive_ptr{this};
+        ON_BLOCK_EXIT([&] {
+            intrusiveThis.detach();
+            unsafeRefDecRefCountTo(0u);
+        });
+        variables.seedVariablesWithLetParameters(intrusiveThis, *letParameters);
+    }
 }
 
 void ExpressionContext::checkForInterrupt() {
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index 0581f7b92d3..80d20ea0054 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -142,6 +142,7 @@ public:
                       std::unique_ptr<CollatorInterface> collator,
                       const NamespaceString& ns,
                       const boost::optional<RuntimeConstants>& runtimeConstants = boost::none,
+                      const boost::optional<BSONObj>& letParameters = boost::none,
                       bool mayDbProfile = true);
 
     /**
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 cdb3a24049e..8f4c9fbb392 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,6 +613,7 @@ Update CommonMongodProcessInterface::buildUpdateOp(
         return wcb;
     }());
     updateOp.setRuntimeConstants(expCtx->getRuntimeConstants());
+    updateOp.setLet(expCtx->variables.serializeLetParameters(expCtx->variablesParseState));
     return updateOp;
 }
 
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index 7d62a437ba1..0ed84ca67be 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -102,6 +102,7 @@ const char QueryRequest::kNoCursorTimeoutField[] = "noCursorTimeout";
 const char QueryRequest::kAwaitDataField[] = "awaitData";
 const char QueryRequest::kPartialResultsField[] = "allowPartialResults";
 const char QueryRequest::kRuntimeConstantsField[] = "runtimeConstants";
+const char QueryRequest::kLetField[] = "let";
 const char QueryRequest::kTermField[] = "term";
 const char QueryRequest::kOptionsField[] = "options";
 const char QueryRequest::kReadOnceField[] = "readOnce";
@@ -353,6 +354,10 @@ StatusWith<unique_ptr<QueryRequest>> QueryRequest::parseFromFindCommand(unique_p
             qr->_runtimeConstants =
                 RuntimeConstants::parse(IDLParserErrorContext(kRuntimeConstantsField),
                                         cmdObj.getObjectField(kRuntimeConstantsField));
+        } else if (fieldName == kLetField) {
+            if (auto status = checkFieldType(el, Object); !status.isOK())
+                return status;
+            qr->_letParameters = el.Obj().getOwned();
         } else if (fieldName == kOptionsField) {
             // 3.0.x versions of the shell may generate an explain of a find command with an
             // 'options' field. We accept this only if the 'options' field is empty so that
@@ -581,6 +586,10 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const {
         rtcBuilder.doneFast();
     }
 
+    if (_letParameters) {
+        cmdBuilder->append(kLetField, *_letParameters);
+    }
+
     if (_replicationTerm) {
         cmdBuilder->append(kTermField, *_replicationTerm);
     }
@@ -1124,6 +1133,9 @@ StatusWith<BSONObj> QueryRequest::asAggregationCommand() const {
         _runtimeConstants->serialize(&rtcBuilder);
         rtcBuilder.doneFast();
     }
+    if (_letParameters) {
+        aggregationBuilder.append(QueryRequest::kLetField, *_letParameters);
+    }
     return StatusWith<BSONObj>(aggregationBuilder.obj());
 }
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h
index 3db2b9b86de..b039c84a03b 100644
--- a/src/mongo/db/query/query_request.h
+++ b/src/mongo/db/query/query_request.h
@@ -72,6 +72,7 @@ public:
     static const char kAwaitDataField[];
     static const char kPartialResultsField[];
     static const char kRuntimeConstantsField[];
+    static const char kLetField[];
     static const char kTermField[];
     static const char kOptionsField[];
     static const char kReadOnceField[];
@@ -367,6 +368,14 @@ public:
         return _runtimeConstants;
     }
 
+    void setLetParameters(BSONObj letParams) {
+        _letParameters = std::move(letParams);
+    }
+
+    const boost::optional<BSONObj>& getLetParameters() const {
+        return _letParameters;
+    }
+
     bool isSlaveOk() const {
         return _slaveOk;
     }
@@ -565,6 +574,10 @@ private:
     // Runtime constants which may be referenced by $expr, if present.
     boost::optional<RuntimeConstants> _runtimeConstants;
 
+    // A document containing user-specified constants. For a find query, these are accessed only
+    // inside $expr.
+    boost::optional<BSONObj> _letParameters;
+
     // Options that can be specified in the OP_QUERY 'flags' header.
     TailableModeEnum _tailableMode = TailableModeEnum::kNormal;
     bool _slaveOk = false;
diff --git a/src/mongo/s/write_ops/batched_command_request.h b/src/mongo/s/write_ops/batched_command_request.h
index dc308c217b7..ee317b1346a 100644
--- a/src/mongo/s/write_ops/batched_command_request.h
+++ b/src/mongo/s/write_ops/batched_command_request.h
@@ -148,6 +148,18 @@ public:
         return _updateReq->getRuntimeConstants();
     }
 
+    void setLet(BSONObj let) {
+        _updateReq->setLet(std::move(let));
+    }
+
+    bool hasLet() const {
+        return _updateReq->getLet().is_initialized();
+    }
+
+    const boost::optional<BSONObj>& getLet() const {
+        return _updateReq->getLet();
+    }
+
     const write_ops::WriteCommandBase& getWriteCommandBase() const;
     void setWriteCommandBase(write_ops::WriteCommandBase writeCommandBase);