SERVER-28773 Create insert mode for SetNode

11 files changed, 99 insertions, 22 deletions

diff --git a/src/mongo/db/update/modifier_table.cpp b/src/mongo/db/update/modifier_table.cpp
index f46b1b1a0d6..72611147891 100644
--- a/src/mongo/db/update/modifier_table.cpp
+++ b/src/mongo/db/update/modifier_table.cpp
@@ -216,6 +216,8 @@ std::unique_ptr<UpdateLeafNode> makeUpdateLeafNode(ModifierType modType) {
             return stdx::make_unique<RenameNode>();
         case MOD_SET:
             return stdx::make_unique<SetNode>();
+        case MOD_SET_ON_INSERT:
+            return stdx::make_unique<SetNode>(UpdateNode::Context::kInsertOnly);
         case MOD_UNSET:
             return stdx::make_unique<UnsetNode>();
         default:
diff --git a/src/mongo/db/update/path_creating_node.cpp b/src/mongo/db/update/path_creating_node.cpp
index 396ffdaf247..920d0054d8a 100644
--- a/src/mongo/db/update/path_creating_node.cpp
+++ b/src/mongo/db/update/path_creating_node.cpp
@@ -108,6 +108,10 @@ void checkImmutablePathsNotModified(mutablebson::Element element,
 }  // namespace
 
 UpdateNode::ApplyResult PathCreatingNode::apply(ApplyParams applyParams) const {
+    if (context == Context::kInsertOnly && !applyParams.insert) {
+        return ApplyResult::noopResult();
+    }
+
     // The value in this Element gets used to create a logging entry (if we have a LogBuilder).
     mutablebson::Element valueToLog = applyParams.element.getDocument().end();
 
diff --git a/src/mongo/db/update/path_creating_node.h b/src/mongo/db/update/path_creating_node.h
index 22c7b6a44ad..2fe69a47064 100644
--- a/src/mongo/db/update/path_creating_node.h
+++ b/src/mongo/db/update/path_creating_node.h
@@ -41,6 +41,8 @@ namespace mongo {
  */
 class PathCreatingNode : public UpdateLeafNode {
 public:
+    explicit PathCreatingNode(Context context = Context::kAll) : UpdateLeafNode(context) {}
+
     ApplyResult apply(ApplyParams applyParams) const final;
 
 protected:
diff --git a/src/mongo/db/update/set_node.h b/src/mongo/db/update/set_node.h
index 626d47be7ed..26c4a66d9a5 100644
--- a/src/mongo/db/update/set_node.h
+++ b/src/mongo/db/update/set_node.h
@@ -38,6 +38,8 @@ namespace mongo {
  */
 class SetNode : public PathCreatingNode {
 public:
+    explicit SetNode(Context context = Context::kAll) : PathCreatingNode(context) {}
+
     Status init(BSONElement modExpr, const CollatorInterface* collator) final;
 
     std::unique_ptr<UpdateNode> clone() const final {
diff --git a/src/mongo/db/update/set_node_test.cpp b/src/mongo/db/update/set_node_test.cpp
index 932798e61e7..3d90d27b6c5 100644
--- a/src/mongo/db/update/set_node_test.cpp
+++ b/src/mongo/db/update/set_node_test.cpp
@@ -1275,5 +1275,58 @@ TEST_F(SetNodeTest, ApplyCanCreatePrefixOfImmutablePath) {
     ASSERT_EQUALS(fromjson("{$set: {a: 2}}"), getLogDoc());
 }
 
+TEST_F(SetNodeTest, ApplySetOnInsertIsNoopWhenInsertIsFalse) {
+    auto update = fromjson("{$setOnInsert: {a: 2}}");
+    const CollatorInterface* collator = nullptr;
+    SetNode node(UpdateNode::Context::kInsertOnly);
+    ASSERT_OK(node.init(update["$setOnInsert"]["a"], collator));
+
+    Document doc(fromjson("{}"));
+    setPathToCreate("a");
+    addIndexedPath("a");
+    auto result = node.apply(getApplyParams(doc.root()));
+    ASSERT_TRUE(result.noop);
+    ASSERT_FALSE(result.indexesAffected);
+    ASSERT_EQUALS(fromjson("{}"), doc);
+    ASSERT_TRUE(doc.isInPlaceModeEnabled());
+    ASSERT_EQUALS(fromjson("{}"), getLogDoc());
+}
+
+TEST_F(SetNodeTest, ApplySetOnInsertIsAppliedWhenInsertIsTrue) {
+    auto update = fromjson("{$setOnInsert: {a: 2}}");
+    const CollatorInterface* collator = nullptr;
+    SetNode node(UpdateNode::Context::kInsertOnly);
+    ASSERT_OK(node.init(update["$setOnInsert"]["a"], collator));
+
+    Document doc(fromjson("{}"));
+    setPathToCreate("a");
+    setInsert(true);
+    addIndexedPath("a");
+    setLogBuilderToNull();  // The log builder is null for inserts.
+    auto result = node.apply(getApplyParams(doc.root()));
+    ASSERT_FALSE(result.noop);
+    ASSERT_TRUE(result.indexesAffected);
+    ASSERT_EQUALS(fromjson("{a: 2}"), doc);
+    ASSERT_FALSE(doc.isInPlaceModeEnabled());
+}
+
+TEST_F(SetNodeTest, ApplySetOnInsertExistingPath) {
+    auto update = fromjson("{$setOnInsert: {a: 2}}");
+    const CollatorInterface* collator = nullptr;
+    SetNode node(UpdateNode::Context::kInsertOnly);
+    ASSERT_OK(node.init(update["$setOnInsert"]["a"], collator));
+
+    Document doc(fromjson("{a: 1}"));
+    setPathTaken("a");
+    setInsert(true);
+    addIndexedPath("a");
+    setLogBuilderToNull();  // The log builder is null for inserts.
+    auto result = node.apply(getApplyParams(doc.root()["a"]));
+    ASSERT_FALSE(result.noop);
+    ASSERT_TRUE(result.indexesAffected);
+    ASSERT_EQUALS(fromjson("{a: 2}"), doc);
+    ASSERT_TRUE(doc.isInPlaceModeEnabled());
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/update/update_driver.cpp b/src/mongo/db/update/update_driver.cpp
index 8402971f9af..477df47b176 100644
--- a/src/mongo/db/update/update_driver.cpp
+++ b/src/mongo/db/update/update_driver.cpp
@@ -291,6 +291,7 @@ Status UpdateDriver::update(StringData matchedField,
         // We parsed using the new UpdateNode implementation.
         UpdateNode::ApplyParams applyParams(doc->root(), immutablePaths);
         applyParams.matchedField = matchedField;
+        applyParams.insert = _insert;
         applyParams.fromReplication = _modOptions.fromReplication;
         applyParams.validateForStorage = validateForStorage;
         applyParams.indexData = _indexedFields;
@@ -320,15 +321,7 @@ Status UpdateDriver::update(StringData matchedField,
                 return status;
             }
 
-            // If a mod wants to be applied only if this is an upsert (or only if this is a
-            // strict update), we should respect that. If a mod doesn't care, it would state
-            // it is fine with ANY update context.
-            const bool validContext =
-                (execInfo.context == ModifierInterface::ExecInfo::ANY_CONTEXT ||
-                 execInfo.context == _context);
-
-            // Nothing to do if not in a valid context.
-            if (!validContext) {
+            if (execInfo.context == ModifierInterface::ExecInfo::INSERT_CONTEXT && !_insert) {
                 continue;
             }
 
@@ -503,14 +496,6 @@ void UpdateDriver::setCollator(const CollatorInterface* collator) {
     _modOptions.collator = collator;
 }
 
-ModifierInterface::ExecInfo::UpdateContext UpdateDriver::context() const {
-    return _context;
-}
-
-void UpdateDriver::setContext(ModifierInterface::ExecInfo::UpdateContext context) {
-    _context = context;
-}
-
 BSONObj UpdateDriver::makeOplogEntryQuery(const BSONObj& doc, bool multi) const {
     BSONObjBuilder idPattern;
     BSONElement id;
diff --git a/src/mongo/db/update/update_driver.h b/src/mongo/db/update/update_driver.h
index c188ae62960..52a9f77e912 100644
--- a/src/mongo/db/update/update_driver.h
+++ b/src/mongo/db/update/update_driver.h
@@ -137,8 +137,9 @@ public:
     ModifierInterface::Options modOptions() const;
     void setModOptions(ModifierInterface::Options modOpts);
 
-    ModifierInterface::ExecInfo::UpdateContext context() const;
-    void setContext(ModifierInterface::ExecInfo::UpdateContext context);
+    void setInsert(bool insert) {
+        _insert = insert;
+    }
 
     mutablebson::Document& getDocument() {
         return _objDoc;
@@ -205,7 +206,7 @@ private:
     bool _positional;
 
     // Is this update going to be an upsert?
-    ModifierInterface::ExecInfo::UpdateContext _context;
+    bool _insert = false;
 
     // The document used to represent or store the object being updated.
     mutablebson::Document _objDoc;
diff --git a/src/mongo/db/update/update_leaf_node.h b/src/mongo/db/update/update_leaf_node.h
index 357def1685e..3e0b5b4d095 100644
--- a/src/mongo/db/update/update_leaf_node.h
+++ b/src/mongo/db/update/update_leaf_node.h
@@ -38,7 +38,7 @@ namespace mongo {
  */
 class UpdateLeafNode : public UpdateNode {
 public:
-    UpdateLeafNode() : UpdateNode(Type::Leaf) {}
+    explicit UpdateLeafNode(Context context = Context::kAll) : UpdateNode(Type::Leaf, context) {}
 
     /**
      * Initializes the leaf node with the value in 'modExpr'. 'modExpr' is a single modifier
diff --git a/src/mongo/db/update/update_node.h b/src/mongo/db/update/update_node.h
index f2dccbad85d..180f5755ecb 100644
--- a/src/mongo/db/update/update_node.h
+++ b/src/mongo/db/update/update_node.h
@@ -59,9 +59,11 @@ class FieldRef;
  */
 class UpdateNode {
 public:
+    enum class Context { kAll, kInsertOnly };
     enum class Type { Object, Array, Leaf, Replacement };
 
-    explicit UpdateNode(Type type) : type(type) {}
+    explicit UpdateNode(Type type, Context context = Context::kAll)
+        : context(context), type(type) {}
     virtual ~UpdateNode() = default;
 
     virtual std::unique_ptr<UpdateNode> clone() const = 0;
@@ -102,6 +104,10 @@ public:
         // index of the array element that caused the query to match the document.
         StringData matchedField;
 
+        // True if the update is being applied to a document to be inserted. $setOnInsert behaves as
+        // a no-op when this flag is false.
+        bool insert = false;
+
         // This is provided because some modifiers may ignore certain errors when the update is from
         // replication.
         bool fromReplication = false;
@@ -151,6 +157,7 @@ public:
                                                                  const UpdateNode& rightNode,
                                                                  FieldRef* pathTaken);
 
+    const Context context;
     const Type type;
 };
 
diff --git a/src/mongo/db/update/update_node_test_fixture.h b/src/mongo/db/update/update_node_test_fixture.h
index 29766992539..b434d93a99e 100644
--- a/src/mongo/db/update/update_node_test_fixture.h
+++ b/src/mongo/db/update/update_node_test_fixture.h
@@ -55,6 +55,7 @@ protected:
         _pathToCreate = std::make_shared<FieldRef>();
         _pathTaken = std::make_shared<FieldRef>();
         _matchedField = StringData();
+        _insert = false;
         _fromReplication = false;
         _validateForStorage = true;
         _indexData.reset();
@@ -67,6 +68,7 @@ protected:
         applyParams.pathToCreate = _pathToCreate;
         applyParams.pathTaken = _pathTaken;
         applyParams.matchedField = _matchedField;
+        applyParams.insert = _insert;
         applyParams.fromReplication = _fromReplication;
         applyParams.validateForStorage = _validateForStorage;
         applyParams.indexData = _indexData.get();
@@ -94,6 +96,10 @@ protected:
         _matchedField = matchedField;
     }
 
+    void setInsert(bool insert) {
+        _insert = insert;
+    }
+
     void setFromReplication(bool fromReplication) {
         _fromReplication = fromReplication;
     }
@@ -123,6 +129,7 @@ private:
     std::shared_ptr<FieldRef> _pathToCreate;
     std::shared_ptr<FieldRef> _pathTaken;
     StringData _matchedField;
+    bool _insert;
     bool _fromReplication;
     bool _validateForStorage;
     std::unique_ptr<UpdateIndexData> _indexData;
diff --git a/src/mongo/db/update/update_object_node_test.cpp b/src/mongo/db/update/update_object_node_test.cpp
index f991d33f7ab..511a38dc461 100644
--- a/src/mongo/db/update/update_object_node_test.cpp
+++ b/src/mongo/db/update/update_object_node_test.cpp
@@ -219,6 +219,20 @@ TEST(UpdateObjectNodeTest, ValidCurrentDatePathParsesSuccessfully) {
                                               foundIdentifiers));
 }
 
+TEST(UpdateObjectNodeTest, ValidSetOnInsertPathParsesSuccessfully) {
+    auto update = fromjson("{$setOnInsert: {'a.b': true}}");
+    const CollatorInterface* collator = nullptr;
+    std::map<StringData, std::unique_ptr<ExpressionWithPlaceholder>> arrayFilters;
+    std::set<std::string> foundIdentifiers;
+    UpdateObjectNode root;
+    ASSERT_OK(UpdateObjectNode::parseAndMerge(&root,
+                                              modifiertable::ModifierType::MOD_SET_ON_INSERT,
+                                              update["$setOnInsert"]["a.b"],
+                                              collator,
+                                              arrayFilters,
+                                              foundIdentifiers));
+}
+
 TEST(UpdateObjectNodeTest, MultiplePositionalElementsFailToParse) {
     auto update = fromjson("{$set: {'a.$.b.$': 5}}");
     const CollatorInterface* collator = nullptr;