SERVER-46715 Support let parameter in findAndModify command

4 files changed, 58 insertions, 1 deletions

diff --git a/jstests/noPassthroughWithMongod/command_let_variables.js b/jstests/noPassthroughWithMongod/command_let_variables.js
index 6f098bde41f..aa04bedbf62 100644
--- a/jstests/noPassthroughWithMongod/command_let_variables.js
+++ b/jstests/noPassthroughWithMongod/command_let_variables.js
@@ -110,8 +110,30 @@ assert.commandFailedWithCode(db.runCommand({
 }),
                              ErrorCodes.TypeMismatch);
 
+// findAndModify
+assert.commandWorked(coll.insert({Species: "spy_bird"}));
+let result = db.runCommand({
+    findAndModify: coll.getName(),
+    let : {target_species: "spy_bird"},
+    query: {$expr: {$eq: ["$Species", "$$target_species"]}},
+    update: {Species: "questionable_bird"},
+    fields: {_id: 0},
+    new: true
+});
+assert.eq(result.value, {Species: "questionable_bird"}, result);
+
+result = db.runCommand({
+    findAndModify: coll.getName(),
+    let : {species_name: "not_a_bird", realSpecies: "dino"},
+    query: {$expr: {$eq: ["$Species", "questionable_bird"]}},
+    update: [{$project: {Species: "$$species_name"}}, {$addFields: {suspect: "$$realSpecies"}}],
+    fields: {_id: 0},
+    new: true
+});
+assert.eq(result.value, {Species: "not_a_bird", suspect: "dino"}, result);
+
 // Delete
-const result = assert.commandWorked(db.runCommand({
+result = assert.commandWorked(db.runCommand({
     delete: coll.getName(),
     let : {target_species: "not_a_bird"},
     deletes: [{q: {$expr: {$eq: ["$Species", "$$target_species"]}}, limit: 0}]
diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index dbb33df7899..80239cc1863 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -129,6 +129,7 @@ void makeUpdateRequest(OperationContext* opCtx,
     requestOut->setUpdateModification(*args.getUpdate());
     requestOut->setRuntimeConstants(
         args.getRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx)));
+    requestOut->setLetParameters(args.getLetParameters());
     requestOut->setSort(args.getSort());
     requestOut->setHint(args.getHint());
     requestOut->setCollation(args.getCollation());
@@ -151,6 +152,7 @@ void makeDeleteRequest(OperationContext* opCtx,
     requestOut->setProj(args.getFields());
     requestOut->setRuntimeConstants(
         args.getRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx)));
+    requestOut->setLet(args.getLetParameters());
     requestOut->setSort(args.getSort());
     requestOut->setHint(args.getHint());
     requestOut->setCollation(args.getCollation());
diff --git a/src/mongo/db/query/find_and_modify_request.cpp b/src/mongo/db/query/find_and_modify_request.cpp
index 703a3621a43..b6d89bdc3b3 100644
--- a/src/mongo/db/query/find_and_modify_request.cpp
+++ b/src/mongo/db/query/find_and_modify_request.cpp
@@ -48,6 +48,7 @@ const char kHintField[] = "hint";
 const char kCollationField[] = "collation";
 const char kArrayFiltersField[] = "arrayFilters";
 const char kRuntimeConstantsField[] = "runtimeConstants";
+const char kLet[] = "let";
 const char kRemoveField[] = "remove";
 const char kUpdateField[] = "update";
 const char kNewField[] = "new";
@@ -135,6 +136,12 @@ BSONObj FindAndModifyRequest::toBSON(const BSONObj& commandPassthroughFields) co
         rtcBuilder.doneFast();
     }
 
+    if (_letParameters) {
+        if (auto letParams = _letParameters.get(); !letParams.isEmpty()) {
+            builder.append(kLet, _letParameters.get());
+        }
+    }
+
     if (_shouldReturnNew) {
         builder.append(kNewField, _shouldReturnNew);
     }
@@ -170,6 +177,7 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
     bool arrayFiltersSet = false;
     std::vector<BSONObj> arrayFilters;
     boost::optional<RuntimeConstants> runtimeConstants;
+    BSONObj letParameters;
     bool writeConcernOptionsSet = false;
     WriteConcernOptions writeConcernOptions;
 
@@ -245,6 +253,14 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
             runtimeConstants =
                 RuntimeConstants::parse(IDLParserErrorContext(kRuntimeConstantsField),
                                         cmdObj.getObjectField(kRuntimeConstantsField));
+        } else if (field == kLet) {
+            BSONElement letElt;
+            if (Status letEltStatus =
+                    bsonExtractTypedField(cmdObj, kLet, BSONType::Object, &letElt);
+                !letEltStatus.isOK()) {
+                return letEltStatus;
+            }
+            letParameters = letElt.embeddedObject();
         } else if (field == kWriteConcernField) {
             BSONElement writeConcernElt;
             Status writeConcernEltStatus = bsonExtractTypedField(
@@ -301,6 +317,7 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
     request.setHint(hint);
     request.setCollation(collation);
     request.setBypassDocumentValidation(bypassDocumentValidation);
+    request.setLetParameters(letParameters);
     if (arrayFiltersSet) {
         request.setArrayFilters(std::move(arrayFilters));
     }
diff --git a/src/mongo/db/query/find_and_modify_request.h b/src/mongo/db/query/find_and_modify_request.h
index e054687e6e6..37249d03d34 100644
--- a/src/mongo/db/query/find_and_modify_request.h
+++ b/src/mongo/db/query/find_and_modify_request.h
@@ -185,6 +185,21 @@ public:
     }
 
     /**
+     * Sets user-specified parameters that can be used by the findAndModify command's pipeline-style
+     * updates and inside $expr.
+     */
+    void setLetParameters(const boost::optional<BSONObj>& letParameters) {
+        _letParameters = letParameters;
+    }
+
+    /**
+     * Returns the runtime constants associated with this findAndModify request, if present.
+     */
+    const boost::optional<BSONObj>& getLetParameters() const {
+        return _letParameters;
+    }
+
+    /**
      * Sets the write concern for this request.
      */
     void setWriteConcern(WriteConcernOptions writeConcern);
@@ -210,6 +225,7 @@ private:
     boost::optional<BSONObj> _collation;
     boost::optional<std::vector<BSONObj>> _arrayFilters;
     boost::optional<RuntimeConstants> _runtimeConstants;
+    boost::optional<BSONObj> _letParameters;
     bool _shouldReturnNew{false};
     boost::optional<WriteConcernOptions> _writeConcern;
     bool _bypassDocumentValidation{false};