SERVER-46718 Support 'let' parameters for findAndModify in sharded environments

13 files changed, 202 insertions, 92 deletions

diff --git a/jstests/core/command_let_variables.js b/jstests/core/command_let_variables.js
index 2ae1b7f5f15..7d7b97d13cd 100644
--- a/jstests/core/command_let_variables.js
+++ b/jstests/core/command_let_variables.js
@@ -308,4 +308,35 @@ assert.commandWorked(testDB.runCommand({
     let : {variable: "Song Thrush"},
     cursor: {}
 }));
+
+// Test that findAndModify works correctly with let parameter arguments.
+assert.commandWorked(coll.insert({_id: 5, Species: "spy_bird"}));
+result = testDB.runCommand({
+    findAndModify: coll.getName(),
+    let : {target_species: "spy_bird"},
+    // Querying on _id field for sharded collection passthroughs.
+    query: {$and: [{_id: 5}, {$expr: {$eq: ["$Species", "$$target_species"]}}]},
+    update: {Species: "questionable_bird"},
+    new: true
+});
+expectedResults = {
+    _id: 5,
+    Species: "questionable_bird"
+};
+assert.eq(expectedResults, result.value, result);
+
+result = testDB.runCommand({
+    findAndModify: coll.getName(),
+    let : {species_name: "not_a_bird", realSpecies: "dino"},
+    // Querying on _id field for sharded collection passthroughs.
+    query: {$and: [{_id: 5}, {$expr: {$eq: ["$Species", "questionable_bird"]}}]},
+    update: [{$project: {Species: "$$species_name"}}, {$addFields: {suspect: "$$realSpecies"}}],
+    new: true
+});
+expectedResults = {
+    _id: 5,
+    Species: "not_a_bird",
+    suspect: "dino"
+};
+assert.eq(expectedResults, result.value, result);
 }());
diff --git a/jstests/noPassthroughWithMongod/command_let_variables.js b/jstests/noPassthroughWithMongod/command_let_variables.js
index 42f5ee4c2fc..7cfc6c1ad5e 100644
--- a/jstests/noPassthroughWithMongod/command_let_variables.js
+++ b/jstests/noPassthroughWithMongod/command_let_variables.js
@@ -196,28 +196,6 @@ assert.commandWorked(db.runCommand({
 }));
 assert.eq(targetColl.aggregate({$match: {$expr: {$eq: ["$var", "INNER"]}}}).toArray().length, 1);
 
-// 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);
-
 // Update
 assert.commandWorked(db.runCommand({
     update: coll.getName(),
diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index ee959b26226..eb8a3f5af60 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -676,6 +676,7 @@ env.Library(
         'db/logical_session_cache',
         'db/logical_session_cache_impl',
         'db/logical_time_metadata_hook',
+        'db/pipeline/process_interface/mongos_process_interface_factory',
         'db/read_write_concern_defaults',
         'db/server_options',
         'db/server_options_base',
diff --git a/src/mongo/db/pipeline/process_interface/SConscript b/src/mongo/db/pipeline/process_interface/SConscript
index 8fc9a8e1fb3..aa783988fd3 100644
--- a/src/mongo/db/pipeline/process_interface/SConscript
+++ b/src/mongo/db/pipeline/process_interface/SConscript
@@ -70,6 +70,16 @@ env.Library(
 )
 
 env.Library(
+    target='mongos_process_interface_factory',
+    source=[
+        'mongos_process_interface_factory.cpp',
+    ],
+    LIBDEPS=[
+        'mongos_process_interface',
+    ],
+)
+
+env.Library(
     target='mongod_process_interface_factory',
     source=[
         'mongod_process_interface_factory.cpp',
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface_factory.cpp b/src/mongo/db/pipeline/process_interface/mongos_process_interface_factory.cpp
new file mode 100644
index 00000000000..88db7a878ae
--- /dev/null
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface_factory.cpp
@@ -0,0 +1,49 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/base/shim.h"
+#include "mongo/db/pipeline/process_interface/mongos_process_interface.h"
+#include "mongo/executor/task_executor_pool.h"
+#include "mongo/s/grid.h"
+
+namespace mongo {
+namespace {
+
+std::shared_ptr<MongoProcessInterface> MongoProcessInterfaceCreateImpl(OperationContext* opCtx) {
+    return std::make_shared<MongosProcessInterface>(
+        Grid::get(opCtx)->getExecutorPool()->getArbitraryExecutor());
+}
+
+auto mongoProcessInterfaceCreateRegistration = MONGO_WEAK_FUNCTION_REGISTRATION(
+    MongoProcessInterface::create, MongoProcessInterfaceCreateImpl);
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index 8afbb78ddb1..5c4fbbc4b34 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -109,7 +109,9 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize(
         // A collator can enter through both the QueryRequest and ExpressionContext arguments.
         // This invariant ensures that both collators are the same because downstream we
         // pull the collator from only one of the ExpressionContext carrier.
-        invariant(CollatorInterface::collatorsMatch(collator.get(), expCtx->getCollator()));
+        if (collator.get() && expCtx->getCollator()) {
+            invariant(CollatorInterface::collatorsMatch(collator.get(), expCtx->getCollator()));
+        }
     }
 
     StatusWithMatchExpression statusWithMatcher = MatchExpressionParser::parse(
diff --git a/src/mongo/s/SConscript b/src/mongo/s/SConscript
index 91e08c4277c..c3ade44368e 100644
--- a/src/mongo/s/SConscript
+++ b/src/mongo/s/SConscript
@@ -46,6 +46,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/commands/txn_cmd_request',
         '$BUILD_DIR/mongo/db/curop',
         '$BUILD_DIR/mongo/db/logical_session_id_helpers',
+        '$BUILD_DIR/mongo/db/pipeline/process_interface/mongo_process_interface',
         '$BUILD_DIR/mongo/db/shared_request_handling',
         '$BUILD_DIR/mongo/db/repl/read_concern_args',
         'async_requests_sender',
diff --git a/src/mongo/s/cluster_commands_helpers.cpp b/src/mongo/s/cluster_commands_helpers.cpp
index 85839a860c8..d992938f47b 100644
--- a/src/mongo/s/cluster_commands_helpers.cpp
+++ b/src/mongo/s/cluster_commands_helpers.cpp
@@ -42,6 +42,7 @@
 #include "mongo/db/logical_clock.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/pipeline/process_interface/mongo_process_interface.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
 #include "mongo/db/query/cursor_response.h"
 #include "mongo/db/repl/read_concern_args.h"
@@ -102,6 +103,46 @@ std::unique_ptr<WriteConcernErrorDetail> getWriteConcernErrorDetailFromBSONObj(c
     return std::make_unique<WriteConcernErrorDetail>(getWriteConcernErrorDetail(wcErrorElem));
 }
 
+boost::intrusive_ptr<ExpressionContext> makeExpressionContextWithDefaultsForTargeter(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const BSONObj& collation,
+    const boost::optional<ExplainOptions::Verbosity>& verbosity,
+    const boost::optional<BSONObj>& letParameters,
+    const boost::optional<RuntimeConstants>& runtimeConstants) {
+
+    auto&& cif = [&]() {
+        if (collation.isEmpty()) {
+            return std::unique_ptr<CollatorInterface>{};
+        } else {
+            return uassertStatusOK(
+                CollatorFactoryInterface::get(opCtx->getServiceContext())->makeFromBSON(collation));
+        }
+    }();
+
+    StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
+    resolvedNamespaces.emplace(nss.coll(),
+                               ExpressionContext::ResolvedNamespace(nss, std::vector<BSONObj>{}));
+
+    return make_intrusive<ExpressionContext>(
+        opCtx,
+        verbosity,
+        true,   // fromMongos
+        false,  // needs merge
+        false,  // disk use is banned on mongos
+        true,   // bypass document validation, mongos isn't a storage node
+        false,  // not mapReduce
+        nss,
+        runtimeConstants,
+        std::move(cif),
+        MongoProcessInterface::create(opCtx),
+        std::move(resolvedNamespaces),
+        boost::none,  // collection uuid
+        letParameters,
+        false  // mongos has no profile collection
+    );
+}
+
 namespace {
 
 const auto kAllowImplicitCollectionCreation = "allowImplicitCollectionCreation"_sd;
diff --git a/src/mongo/s/cluster_commands_helpers.h b/src/mongo/s/cluster_commands_helpers.h
index 096a04019e1..b66f3a2f95f 100644
--- a/src/mongo/s/cluster_commands_helpers.h
+++ b/src/mongo/s/cluster_commands_helpers.h
@@ -65,6 +65,18 @@ void appendWriteConcernErrorToCmdResponse(const ShardId& shardID,
 std::unique_ptr<WriteConcernErrorDetail> getWriteConcernErrorDetailFromBSONObj(const BSONObj& obj);
 
 /**
+ * Makes an expression context suitable for canonicalization of queries that contain let parameters
+ * and runtimeConstants on mongos.
+ */
+boost::intrusive_ptr<ExpressionContext> makeExpressionContextWithDefaultsForTargeter(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const BSONObj& collation,
+    const boost::optional<ExplainOptions::Verbosity>& verbosity,
+    const boost::optional<BSONObj>& letParameters,
+    const boost::optional<RuntimeConstants>& runtimeConstants);
+
+/**
  * Consults the routing info to build requests for:
  * 1. If sharded, shards that own chunks for the namespace, or
  * 2. If unsharded, the primary shard for the database.
diff --git a/src/mongo/s/commands/SConscript b/src/mongo/s/commands/SConscript
index 3db993ea17e..b11a0394cf2 100644
--- a/src/mongo/s/commands/SConscript
+++ b/src/mongo/s/commands/SConscript
@@ -174,5 +174,6 @@ env.CppUnitTest(
         'cluster_command_test_fixture',
         '$BUILD_DIR/mongo/db/auth/authmocks',
         '$BUILD_DIR/mongo/db/auth/saslauth',
+        '$BUILD_DIR/mongo/db/pipeline/process_interface/mongos_process_interface_factory',
     ],
 )
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index 8a8da3bfb44..5a388477214 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -81,12 +81,36 @@ BSONObj getCollation(const BSONObj& cmdObj) {
     return BSONObj();
 }
 
+boost::optional<BSONObj> getLet(const BSONObj& cmdObj) {
+    if (auto letElem = cmdObj.getField("let"_sd); letElem.type() == BSONType::Object) {
+        auto bob = BSONObjBuilder();
+        bob.appendElementsUnique(letElem.embeddedObject());
+        return bob.obj();
+    }
+    return boost::none;
+}
+
+boost::optional<RuntimeConstants> getRuntimeConstants(const BSONObj& cmdObj) {
+    if (auto rcElem = cmdObj.getField("runtimeConstants"_sd); rcElem.type() == BSONType::Object) {
+        IDLParserErrorContext ctx("internalRuntimeConstants");
+        return RuntimeConstants::parse(ctx, rcElem.embeddedObject());
+    }
+    return boost::none;
+}
+
 BSONObj getShardKey(OperationContext* opCtx,
                     const ChunkManager& chunkMgr,
                     const NamespaceString& nss,
-                    const BSONObj& query) {
+                    const BSONObj& query,
+                    const BSONObj& collation,
+                    const boost::optional<ExplainOptions::Verbosity> verbosity,
+                    const boost::optional<BSONObj>& let,
+                    const boost::optional<RuntimeConstants>& runtimeConstants) {
+    auto expCtx = makeExpressionContextWithDefaultsForTargeter(
+        opCtx, nss, collation, verbosity, let, runtimeConstants);
+
     BSONObj shardKey =
-        uassertStatusOK(chunkMgr.getShardKeyPattern().extractShardKeyFromQuery(opCtx, nss, query));
+        uassertStatusOK(chunkMgr.getShardKeyPattern().extractShardKeyFromQuery(expCtx, query));
     uassert(ErrorCodes::ShardKeyNotFound,
             "Query for sharded findAndModify must contain the shard key",
             !shardKey.isEmpty());
@@ -191,7 +215,10 @@ public:
 
             const BSONObj query = cmdObj.getObjectField("query");
             const BSONObj collation = getCollation(cmdObj);
-            const BSONObj shardKey = getShardKey(opCtx, *chunkMgr, nss, query);
+            const auto let = getLet(cmdObj);
+            const auto rc = getRuntimeConstants(cmdObj);
+            const BSONObj shardKey =
+                getShardKey(opCtx, *chunkMgr, nss, query, collation, verbosity, let, rc);
             const auto chunk = chunkMgr->findIntersectingChunk(shardKey, collation);
 
             shard = uassertStatusOK(
@@ -266,7 +293,10 @@ public:
 
         const BSONObj query = cmdObjForShard.getObjectField("query");
         const BSONObj collation = getCollation(cmdObjForShard);
-        const BSONObj shardKey = getShardKey(opCtx, *chunkMgr, nss, query);
+        const auto let = getLet(cmdObjForShard);
+        const auto rc = getRuntimeConstants(cmdObjForShard);
+        const BSONObj shardKey =
+            getShardKey(opCtx, *chunkMgr, nss, query, collation, boost::none, let, rc);
         auto chunk = chunkMgr->findIntersectingChunk(shardKey, collation);
 
         _runCommand(opCtx,
diff --git a/src/mongo/s/query/cluster_aggregate.cpp b/src/mongo/s/query/cluster_aggregate.cpp
index 2dfebe0aab0..2864f83f4dd 100644
--- a/src/mongo/s/query/cluster_aggregate.cpp
+++ b/src/mongo/s/query/cluster_aggregate.cpp
@@ -116,14 +116,14 @@ boost::intrusive_ptr<ExpressionContext> makeExpressionContext(
 
     // Create the expression context, and set 'inMongos' to true. We explicitly do *not* set
     // mergeCtx->tempDir.
-    auto mergeCtx =
-        new ExpressionContext(opCtx,
-                              request,
-                              std::move(collation),
-                              std::make_shared<MongosProcessInterface>(
-                                  Grid::get(opCtx)->getExecutorPool()->getArbitraryExecutor()),
-                              std::move(resolvedNamespaces),
-                              uuid);
+    auto mergeCtx = make_intrusive<ExpressionContext>(
+        opCtx,
+        request,
+        std::move(collation),
+        std::make_shared<MongosProcessInterface>(
+            Grid::get(opCtx)->getExecutorPool()->getArbitraryExecutor()),
+        std::move(resolvedNamespaces),
+        uuid);
 
     mergeCtx->inMongos = true;
     return mergeCtx;
diff --git a/src/mongo/s/write_ops/chunk_manager_targeter.cpp b/src/mongo/s/write_ops/chunk_manager_targeter.cpp
index e6a20e2786c..ae40d648c28 100644
--- a/src/mongo/s/write_ops/chunk_manager_targeter.cpp
+++ b/src/mongo/s/write_ops/chunk_manager_targeter.cpp
@@ -119,9 +119,8 @@ UpdateType getUpdateExprType(const write_ops::UpdateOpEntry& updateDoc) {
  * generating the new document in the case of an upsert. It is therefore always correct to target
  * the operation on the basis of the combined updateExpr and query.
  */
-BSONObj getUpdateExprForTargeting(OperationContext* opCtx,
+BSONObj getUpdateExprForTargeting(const boost::intrusive_ptr<ExpressionContext> expCtx,
                                   const ShardKeyPattern& shardKeyPattern,
-                                  const NamespaceString& nss,
                                   UpdateType updateType,
                                   const write_ops::UpdateOpEntry& updateOp) {
     // We should never see an invalid update type here.
@@ -152,7 +151,7 @@ BSONObj getUpdateExprForTargeting(OperationContext* opCtx,
     // This will guarantee that we can target a single shard, but it is not necessarily fatal if no
     // exact _id can be found.
     const auto idFromQuery =
-        uassertStatusOK(kVirtualIdShardKey.extractShardKeyFromQuery(opCtx, nss, updateOp.getQ()));
+        uassertStatusOK(kVirtualIdShardKey.extractShardKeyFromQuery(expCtx, updateOp.getQ()));
     if (auto idElt = idFromQuery[kIdFieldName]) {
         updateExpr = updateExpr.addField(idElt);
     }
@@ -353,52 +352,6 @@ bool wasMetadataRefreshed(const std::shared_ptr<ChunkManager>& managerA,
 
     return false;
 }
-
-/**
- * Makes an expression context suitable for canonicalization of queries that contain let parameters
- * and runtimeConstants on mongos.
- */
-auto makeExpressionContextWithDefaultsForTargeter(
-    OperationContext* opCtx,
-    const NamespaceString& nss,
-    const BSONObj& collation,
-    const boost::optional<ExplainOptions::Verbosity>& verbosity,
-    const boost::optional<BSONObj>& letParameters,
-    const boost::optional<RuntimeConstants>& runtimeConstants) {
-
-    auto&& cif = [&]() {
-        if (collation.isEmpty()) {
-            return std::unique_ptr<CollatorInterface>{};
-        } else {
-            return uassertStatusOK(
-                CollatorFactoryInterface::get(opCtx->getServiceContext())->makeFromBSON(collation));
-        }
-    }();
-
-    StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
-    resolvedNamespaces.emplace(nss.coll(),
-                               ExpressionContext::ResolvedNamespace(nss, std::vector<BSONObj>{}));
-
-    return make_intrusive<ExpressionContext>(
-        opCtx,
-        verbosity,
-        true,   // fromMongos
-        false,  // needs merge
-        false,  // disk use is banned on mongos
-        true,   // bypass document validation, mongos isn't a storage node
-        false,  // not mapReduce
-        nss,
-        runtimeConstants,
-        std::move(cif),
-        std::make_shared<MongosProcessInterface>(
-            Grid::get(opCtx)->getExecutorPool()->getArbitraryExecutor()),
-        std::move(resolvedNamespaces),
-        boost::none,  // collection uuid
-        letParameters,
-        false  // mongos has no profile collection
-    );
-}
-
 }  // namespace
 
 ChunkManagerTargeter::ChunkManagerTargeter(OperationContext* opCtx,
@@ -478,8 +431,15 @@ std::vector<ShardEndpoint> ChunkManagerTargeter::targetUpdate(OperationContext*
     const auto& shardKeyPattern = _routingInfo->cm()->getShardKeyPattern();
     const auto collation = write_ops::collationOf(updateOp);
 
+    auto expCtx = makeExpressionContextWithDefaultsForTargeter(opCtx,
+                                                               _nss,
+                                                               collation,
+                                                               boost::none,  // explain
+                                                               itemRef.getLet(),
+                                                               itemRef.getRuntimeConstants());
+
     const auto updateExpr =
-        getUpdateExprForTargeting(opCtx, shardKeyPattern, _nss, updateType, updateOp);
+        getUpdateExprForTargeting(expCtx, shardKeyPattern, updateType, updateOp);
     const bool isUpsert = updateOp.getUpsert();
     const auto query = updateOp.getQ();
 
@@ -497,18 +457,12 @@ std::vector<ShardEndpoint> ChunkManagerTargeter::targetUpdate(OperationContext*
     // to target based on the replacement doc, it could result in an insertion even if a document
     // matching the query exists on another shard.
     if (isUpsert) {
-        return targetByShardKey(shardKeyPattern.extractShardKeyFromQuery(opCtx, _nss, query),
+        return targetByShardKey(shardKeyPattern.extractShardKeyFromQuery(expCtx, query),
                                 "Failed to target upsert by query");
     }
 
     // We first try to target based on the update's query. It is always valid to forward any update
     // or upsert to a single shard, so return immediately if we are able to target a single shard.
-    auto expCtx = makeExpressionContextWithDefaultsForTargeter(opCtx,
-                                                               _nss,
-                                                               collation,
-                                                               boost::none,  // explain
-                                                               itemRef.getLet(),
-                                                               itemRef.getRuntimeConstants());
     auto endPoints = uassertStatusOK(_targetQuery(expCtx, query, collation));
     if (endPoints.size() == 1) {
         return endPoints;