SERVER-40397 Add pipeline updates to findAndModify

Adds the ability to specify an array representing an aggregation
pipeline to the 'udpate' argument of findAndModify. Certain options like
'sort' and 'fields' are not yet supported.

8 files changed, 155 insertions, 62 deletions

diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index aa091e04919..5620075aa3c 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -110,7 +110,8 @@ void makeUpdateRequest(const OperationContext* opCtx,
                        UpdateRequest* requestOut) {
     requestOut->setQuery(args.getQuery());
     requestOut->setProj(args.getFields());
-    requestOut->setUpdateModification(args.getUpdateObj());
+    invariant(args.getUpdate());
+    requestOut->setUpdateModification(*args.getUpdate());
     requestOut->setSort(args.getSort());
     requestOut->setCollation(args.getCollation());
     requestOut->setArrayFilters(args.getArrayFilters());
@@ -213,7 +214,11 @@ public:
     }
 
     bool supportsWriteConcern(const BSONObj& cmd) const override {
-        return true;
+        const auto request(uassertStatusOK(FindAndModifyRequest::parseFromBSON(
+            CommandHelpers::parseNsCollectionRequired("test", cmd), cmd)));
+        // TODO SERVER-40403 Add testing for write concern.
+        return request.isRemove() ||
+            request.getUpdate()->type() == write_ops::UpdateModification::Type::kClassic;
     }
 
     void addRequiredPrivileges(const std::string& dbname,
@@ -304,8 +309,14 @@ public:
         OpDebug* const opDebug = &curOp->debug();
 
         boost::optional<DisableDocumentValidation> maybeDisableValidation;
-        if (shouldBypassDocumentValidationForCommand(cmdObj))
+        if (shouldBypassDocumentValidationForCommand(cmdObj)) {
+            // TODO SERVER-40401 Add support for bypassDocumentValidation.
+            uassert(ErrorCodes::NotImplemented,
+                    "No support for pipeline style updates and bypassDocumentValidation",
+                    !args.getUpdate() ||
+                        args.getUpdate()->type() != write_ops::UpdateModification::Type::kPipeline);
             maybeDisableValidation.emplace(opCtx);
+        }
 
         const auto txnParticipant = TransactionParticipant::get(opCtx);
         const auto inTransaction = txnParticipant && txnParticipant.inMultiDocumentTransaction();
diff --git a/src/mongo/db/ops/SConscript b/src/mongo/db/ops/SConscript
index 727eed3a62f..9e679932d08 100644
--- a/src/mongo/db/ops/SConscript
+++ b/src/mongo/db/ops/SConscript
@@ -34,8 +34,8 @@ env.Library(
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/commands/test_commands_enabled',
         '$BUILD_DIR/mongo/db/dbmessage',
-        '$BUILD_DIR/mongo/db/pipeline/aggregation_request',
         '$BUILD_DIR/mongo/idl/idl_parser',
     ],
 )
diff --git a/src/mongo/db/pipeline/aggregation_request.cpp b/src/mongo/db/pipeline/aggregation_request.cpp
index 2357693d4f0..6cf7b38c573 100644
--- a/src/mongo/db/pipeline/aggregation_request.cpp
+++ b/src/mongo/db/pipeline/aggregation_request.cpp
@@ -64,24 +64,6 @@ constexpr StringData AggregationRequest::kExchangeName;
 
 constexpr long long AggregationRequest::kDefaultBatchSize;
 
-StatusWith<std::vector<BSONObj>> AggregationRequest::parsePipelineFromBSON(
-    BSONElement pipelineElem) {
-    std::vector<BSONObj> pipeline;
-    if (pipelineElem.eoo() || pipelineElem.type() != BSONType::Array) {
-        return {ErrorCodes::TypeMismatch, "'pipeline' option must be specified as an array"};
-    }
-
-    for (auto elem : pipelineElem.Obj()) {
-        if (elem.type() != BSONType::Object) {
-            return {ErrorCodes::TypeMismatch,
-                    "Each element of the 'pipeline' array must be an object"};
-        }
-        pipeline.push_back(elem.embeddedObject().getOwned());
-    }
-
-    return std::move(pipeline);
-}
-
 StatusWith<AggregationRequest> AggregationRequest::parseFromBSON(
     const std::string& dbName,
     const BSONObj& cmdObj,
diff --git a/src/mongo/db/pipeline/aggregation_request.h b/src/mongo/db/pipeline/aggregation_request.h
index fc146ef4324..944423ba210 100644
--- a/src/mongo/db/pipeline/aggregation_request.h
+++ b/src/mongo/db/pipeline/aggregation_request.h
@@ -72,7 +72,23 @@ public:
      * Parse an aggregation pipeline definition from 'pipelineElem'. Returns a non-OK status if
      * pipeline is not an array or if any of the array elements are not objects.
      */
-    static StatusWith<std::vector<BSONObj>> parsePipelineFromBSON(BSONElement pipelineElem);
+    static StatusWith<std::vector<BSONObj>> parsePipelineFromBSON(BSONElement pipelineElem) {
+        std::vector<BSONObj> pipeline;
+        if (pipelineElem.eoo() || pipelineElem.type() != BSONType::Array) {
+            return {ErrorCodes::TypeMismatch, "'pipeline' option must be specified as an array"};
+        }
+
+        for (auto elem : pipelineElem.Obj()) {
+            if (elem.type() != BSONType::Object) {
+                return {ErrorCodes::TypeMismatch,
+                        "Each element of the 'pipeline' array must be an object"};
+            }
+            pipeline.push_back(elem.embeddedObject().getOwned());
+        }
+
+        return std::move(pipeline);
+    }
+
 
     /**
      * Create a new instance of AggregationRequest by parsing the raw command object. Returns a
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index 476f9cf4570..7a47464da57 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -173,6 +173,7 @@ env.Library(
         '$BUILD_DIR/mongo/bson/util/bson_extract',
         '$BUILD_DIR/mongo/db/common',
         '$BUILD_DIR/mongo/db/namespace_string',
+        '$BUILD_DIR/mongo/db/ops/write_ops_parsers',
         '$BUILD_DIR/mongo/db/repl/optime',
         "$BUILD_DIR/mongo/idl/idl_parser",
         '$BUILD_DIR/mongo/rpc/command_status',
diff --git a/src/mongo/db/query/find_and_modify_request.cpp b/src/mongo/db/query/find_and_modify_request.cpp
index 6c46b639177..1c5c4961569 100644
--- a/src/mongo/db/query/find_and_modify_request.cpp
+++ b/src/mongo/db/query/find_and_modify_request.cpp
@@ -62,21 +62,19 @@ const std::vector<StringData> _knownFields{
 };
 }  // unnamed namespace
 
-FindAndModifyRequest::FindAndModifyRequest(NamespaceString fullNs, BSONObj query, BSONObj updateObj)
-    : _ns(std::move(fullNs)),
-      _query(query.getOwned()),
-      _updateObj(updateObj.getOwned()),
-      _isRemove(false) {}
+FindAndModifyRequest::FindAndModifyRequest(NamespaceString fullNs,
+                                           BSONObj query,
+                                           boost::optional<write_ops::UpdateModification> update)
+    : _ns(std::move(fullNs)), _query(query.getOwned()), _update(std::move(update)) {}
 
 FindAndModifyRequest FindAndModifyRequest::makeUpdate(NamespaceString fullNs,
                                                       BSONObj query,
-                                                      BSONObj updateObj) {
-    return FindAndModifyRequest(fullNs, query, updateObj);
+                                                      write_ops::UpdateModification update) {
+    return FindAndModifyRequest(fullNs, query, std::move(update));
 }
 
 FindAndModifyRequest FindAndModifyRequest::makeRemove(NamespaceString fullNs, BSONObj query) {
     FindAndModifyRequest request(fullNs, query, {});
-    request._isRemove = true;
     return request;
 }
 
@@ -86,14 +84,14 @@ BSONObj FindAndModifyRequest::toBSON(const BSONObj& commandPassthroughFields) co
     builder.append(kCmdName, _ns.coll());
     builder.append(kQueryField, _query);
 
-    if (_isRemove) {
-        builder.append(kRemoveField, true);
-    } else {
-        builder.append(kUpdateField, _updateObj);
+    if (_update) {
+        _update->serializeToBSON(kUpdateField, &builder);
 
         if (_isUpsert) {
             builder.append(kUpsertField, _isUpsert.get());
         }
+    } else {
+        builder.append(kRemoveField, true);
     }
 
     if (_fieldProjection) {
@@ -136,11 +134,12 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
     BSONObj fields;
     BSONObj updateObj;
     BSONObj sort;
+    boost::optional<write_ops::UpdateModification> update;
+
     BSONObj collation;
     bool shouldReturnNew = false;
     bool isUpsert = false;
     bool isRemove = false;
-    bool isUpdate = false;
     bool arrayFiltersSet = false;
     std::vector<BSONObj> arrayFilters;
 
@@ -152,8 +151,7 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
         } else if (field == kRemoveField) {
             isRemove = cmdObj[kRemoveField].trueValue();
         } else if (field == kUpdateField) {
-            updateObj = cmdObj.getObjectField(kUpdateField);
-            isUpdate = true;
+            update = write_ops::UpdateModification::parseFromBSON(cmdObj[kUpdateField]);
         } else if (field == kNewField) {
             shouldReturnNew = cmdObj[kNewField].trueValue();
         } else if (field == kFieldProjectionField) {
@@ -195,12 +193,12 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
         }
     }
 
-    if (!isRemove && !isUpdate) {
+    if (!isRemove && !update) {
         return {ErrorCodes::FailedToParse, "Either an update or remove=true must be specified"};
     }
 
     if (isRemove) {
-        if (isUpdate) {
+        if (update) {
             return {ErrorCodes::FailedToParse, "Cannot specify both an update and remove=true"};
         }
 
@@ -219,12 +217,31 @@ StatusWith<FindAndModifyRequest> FindAndModifyRequest::parseFromBSON(NamespaceSt
         }
     }
 
-    FindAndModifyRequest request(std::move(fullNs), query, updateObj);
-    request._isRemove = isRemove;
+    if (update && update->type() == write_ops::UpdateModification::Type::kPipeline) {
+        if (arrayFiltersSet) {
+            return {ErrorCodes::FailedToParse, "Cannot specify arrayFilters and a pipeline update"};
+        }
+        if (!collation.isEmpty()) {
+            // TODO SERVER-40399
+            return {ErrorCodes::NotImplemented, "No support for collation and pipeline update"};
+        }
+        if (!fields.isEmpty()) {
+            // TODO SERVER-40404
+            return {ErrorCodes::NotImplemented, "No support for fields and pipeline update"};
+        }
+        if (!sort.isEmpty()) {
+            // TODO SERVER-40405
+            return {ErrorCodes::NotImplemented, "No support for sort and pipeline update"};
+        }
+    }
+
+    FindAndModifyRequest request(std::move(fullNs), query, std::move(update));
     request.setFieldProjection(fields);
     request.setSort(sort);
     request.setCollation(collation);
-    request.setArrayFilters(std::move(arrayFilters));
+    if (arrayFiltersSet) {
+        request.setArrayFilters(std::move(arrayFilters));
+    }
 
     if (!isRemove) {
         request.setShouldReturnNew(shouldReturnNew);
@@ -257,16 +274,16 @@ void FindAndModifyRequest::setQuery(BSONObj query) {
     _query = query.getOwned();
 }
 void FindAndModifyRequest::setUpdateObj(BSONObj updateObj) {
-    _updateObj = updateObj.getOwned();
+    _update.emplace(updateObj.getOwned());
 }
 
 void FindAndModifyRequest::setShouldReturnNew(bool shouldReturnNew) {
-    dassert(!_isRemove);
+    dassert(_update);
     _shouldReturnNew = shouldReturnNew;
 }
 
 void FindAndModifyRequest::setUpsert(bool upsert) {
-    dassert(!_isRemove);
+    dassert(_update);
     _isUpsert = upsert;
 }
 
@@ -286,8 +303,8 @@ BSONObj FindAndModifyRequest::getFields() const {
     return _fieldProjection.value_or(BSONObj());
 }
 
-BSONObj FindAndModifyRequest::getUpdateObj() const {
-    return _updateObj;
+const boost::optional<write_ops::UpdateModification>& FindAndModifyRequest::getUpdate() const {
+    return _update;
 }
 
 BSONObj FindAndModifyRequest::getSort() const {
@@ -314,6 +331,6 @@ bool FindAndModifyRequest::isUpsert() const {
 }
 
 bool FindAndModifyRequest::isRemove() const {
-    return _isRemove;
+    return !static_cast<bool>(_update);
 }
 }
diff --git a/src/mongo/db/query/find_and_modify_request.h b/src/mongo/db/query/find_and_modify_request.h
index f0882a63941..897bce27c06 100644
--- a/src/mongo/db/query/find_and_modify_request.h
+++ b/src/mongo/db/query/find_and_modify_request.h
@@ -33,6 +33,7 @@
 
 #include "mongo/db/jsobj.h"
 #include "mongo/db/namespace_string.h"
+#include "mongo/db/ops/write_ops_parsers.h"
 #include "mongo/db/write_concern_options.h"
 
 namespace mongo {
@@ -55,11 +56,11 @@ public:
     static constexpr auto kCommandName = "findAndModify"_sd;
 
     /**
-     * Creates a new instance of an 'update' type findAndModify request.
+     * Creates a new instance of a 'update' type findAndModify request.
      */
     static FindAndModifyRequest makeUpdate(NamespaceString fullNs,
                                            BSONObj query,
-                                           BSONObj updateObj);
+                                           write_ops::UpdateModification update);
 
     /**
      * Creates a new instance of an 'remove' type findAndModify request.
@@ -99,7 +100,7 @@ public:
     const NamespaceString& getNamespaceString() const;
     BSONObj getQuery() const;
     BSONObj getFields() const;
-    BSONObj getUpdateObj() const;
+    const boost::optional<write_ops::UpdateModification>& getUpdate() const;
     BSONObj getSort() const;
     BSONObj getCollation() const;
     const std::vector<BSONObj>& getArrayFilters() const;
@@ -171,15 +172,14 @@ private:
     /**
      * Creates a new FindAndModifyRequest with the required fields.
      */
-    FindAndModifyRequest(NamespaceString fullNs, BSONObj query, BSONObj updateObj);
+    FindAndModifyRequest(NamespaceString fullNs,
+                         BSONObj query,
+                         boost::optional<write_ops::UpdateModification> update);
 
     // Required fields
     const NamespaceString _ns;
     BSONObj _query;
 
-    // Required for updates
-    BSONObj _updateObj;
-
     boost::optional<bool> _isUpsert;
     boost::optional<BSONObj> _fieldProjection;
     boost::optional<BSONObj> _sort;
@@ -188,7 +188,7 @@ private:
     boost::optional<bool> _shouldReturnNew;
     boost::optional<WriteConcernOptions> _writeConcern;
 
-    // Flag used internally to differentiate whether this is an update or remove type request.
-    bool _isRemove;
+    // Holds value when performing an update request and none when a remove request.
+    boost::optional<write_ops::UpdateModification> _update;
 };
 }
diff --git a/src/mongo/db/query/find_and_modify_request_test.cpp b/src/mongo/db/query/find_and_modify_request_test.cpp
index a30bdd2e538..d50d456f147 100644
--- a/src/mongo/db/query/find_and_modify_request_test.cpp
+++ b/src/mongo/db/query/find_and_modify_request_test.cpp
@@ -30,6 +30,7 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/bson/json.h"
+#include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/query/find_and_modify_request.h"
 #include "mongo/unittest/unittest.h"
 
@@ -50,6 +51,23 @@ TEST(FindAndModifyRequest, BasicUpdate) {
     ASSERT_BSONOBJ_EQ(expectedObj, request.toBSON({}));
 }
 
+TEST(FindAndModifyRequest, PipelineUpdate) {
+    setTestCommandsEnabled(true);
+    const BSONObj query(BSON("x" << 1));
+    const BSONObj pipelineBSON(
+        BSON("pipeline" << BSON_ARRAY(BSON("$addFields" << BSON("y" << 1)))));
+    auto request = FindAndModifyRequest::makeUpdate(
+        NamespaceString("test.user"), query, pipelineBSON["pipeline"]);
+
+    BSONObj expectedObj(fromjson(R"json({
+            findAndModify: 'user',
+            query: { x: 1 },
+            update: [{$addFields: {y: 1}}]
+        })json"));
+
+    ASSERT_BSONOBJ_EQ(expectedObj, request.toBSON({}));
+}
+
 TEST(FindAndModifyRequest, UpdateWithUpsert) {
     const BSONObj query(BSON("x" << 1));
     const BSONObj update(BSON("y" << 1));
@@ -363,7 +381,9 @@ TEST(FindAndModifyRequest, ParseWithUpdateOnlyRequiredFields) {
     auto request = parseStatus.getValue();
     ASSERT_EQUALS(NamespaceString("a.b").toString(), request.getNamespaceString().toString());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1), request.getQuery());
-    ASSERT_BSONOBJ_EQ(BSON("y" << 1), request.getUpdateObj());
+    ASSERT(request.getUpdate());
+    ASSERT(request.getUpdate()->type() == write_ops::UpdateModification::Type::kClassic);
+    ASSERT_BSONOBJ_EQ(BSON("y" << 1), request.getUpdate()->getUpdateClassic());
     ASSERT_EQUALS(false, request.isUpsert());
     ASSERT_EQUALS(false, request.isRemove());
     ASSERT_BSONOBJ_EQ(BSONObj(), request.getFields());
@@ -391,7 +411,9 @@ TEST(FindAndModifyRequest, ParseWithUpdateFullSpec) {
     auto request = parseStatus.getValue();
     ASSERT_EQUALS(NamespaceString("a.b").toString(), request.getNamespaceString().toString());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1), request.getQuery());
-    ASSERT_BSONOBJ_EQ(BSON("y" << 1), request.getUpdateObj());
+    ASSERT(request.getUpdate());
+    ASSERT(request.getUpdate()->type() == write_ops::UpdateModification::Type::kClassic);
+    ASSERT_BSONOBJ_EQ(BSON("y" << 1), request.getUpdate()->getUpdateClassic());
     ASSERT_EQUALS(true, request.isUpsert());
     ASSERT_EQUALS(false, request.isRemove());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1 << "y" << 1), request.getFields());
@@ -416,7 +438,7 @@ TEST(FindAndModifyRequest, ParseWithRemoveOnlyRequiredFields) {
     auto request = parseStatus.getValue();
     ASSERT_EQUALS(NamespaceString("a.b").toString(), request.getNamespaceString().toString());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1), request.getQuery());
-    ASSERT_BSONOBJ_EQ(BSONObj(), request.getUpdateObj());
+    ASSERT_FALSE(request.getUpdate());
     ASSERT_EQUALS(false, request.isUpsert());
     ASSERT_EQUALS(true, request.isRemove());
     ASSERT_BSONOBJ_EQ(BSONObj(), request.getFields());
@@ -441,7 +463,7 @@ TEST(FindAndModifyRequest, ParseWithRemoveFullSpec) {
     auto request = parseStatus.getValue();
     ASSERT_EQUALS(NamespaceString("a.b").toString(), request.getNamespaceString().toString());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1), request.getQuery());
-    ASSERT_BSONOBJ_EQ(BSONObj(), request.getUpdateObj());
+    ASSERT_FALSE(request.getUpdate());
     ASSERT_EQUALS(false, request.isUpsert());
     ASSERT_EQUALS(true, request.isRemove());
     ASSERT_BSONOBJ_EQ(BSON("x" << 1 << "y" << 1), request.getFields());
@@ -521,5 +543,49 @@ TEST(FindAndModifyRequest, ParseWithCollationTypeMismatch) {
     ASSERT_EQUALS(parseStatus.getStatus(), ErrorCodes::TypeMismatch);
 }
 
+TEST(FindAndModifyRequest, ParsesAndSerializesPipelineUpdate) {
+    setTestCommandsEnabled(true);
+    BSONObj cmdObj(fromjson(R"json({
+            query: { x: 1 },
+            update: [{$replaceRoot: {newRoot: {y: 1}}}]
+        })json"));
+
+    auto request =
+        unittest::assertGet(FindAndModifyRequest::parseFromBSON(NamespaceString("a.b"), cmdObj));
+    ASSERT(request.getUpdate());
+    ASSERT(request.getUpdate()->type() == write_ops::UpdateModification::Type::kPipeline);
+    auto serialized = request.toBSON({});
+    ASSERT_BSONOBJ_EQ(serialized, fromjson(R"json({
+      findAndModify: "b",
+      query: {x: 1},
+      update: [{$replaceRoot: {newRoot: {y: 1}}}],
+      upsert: false,
+      fields: {},
+      sort: {},
+      collation: {},
+      new: false
+    })json"));
+    ASSERT_OK(FindAndModifyRequest::parseFromBSON(NamespaceString("a.b"), serialized).getStatus());
+}
+
+TEST(FindAndModifyRequest, RejectsBothArrayFiltersAndPipelineUpdate) {
+    setTestCommandsEnabled(true);
+    BSONObj cmdObj(fromjson(R"json({
+            query: { x: 1 },
+            update: [{$replaceRoot: {newRoot: {y: 1}}}],
+            arrayFilters: []
+        })json"));
+
+    auto swRequestNoFilters = FindAndModifyRequest::parseFromBSON(NamespaceString("a.b"), cmdObj);
+    ASSERT_EQ(swRequestNoFilters.getStatus(), ErrorCodes::FailedToParse);
+
+    cmdObj = fromjson(R"json({
+            query: { x: 1 },
+            update: [{$replaceRoot: {newRoot: {y: 1}}}],
+            arrayFilters: [{"i.x": 1}]
+        })json");
+    auto swRequestOneFilter = FindAndModifyRequest::parseFromBSON(NamespaceString("a.b"), cmdObj);
+    ASSERT_EQ(swRequestOneFilter.getStatus(), ErrorCodes::FailedToParse);
+}
 }  // unnamed namespace
 }  // namespace mongo