SERVER-39060 Add upsert function to Stitch Support Library

7 files changed, 39 insertions, 20 deletions

diff --git a/src/mongo/db/auth/authz_manager_external_state_mock.cpp b/src/mongo/db/auth/authz_manager_external_state_mock.cpp
index a09fcb9fc16..beda9770230 100644
--- a/src/mongo/db/auth/authz_manager_external_state_mock.cpp
+++ b/src/mongo/db/auth/authz_manager_external_state_mock.cpp
@@ -198,9 +198,10 @@ Status AuthzManagerExternalStateMock::updateOne(OperationContext* opCtx,
         document.reset(*iter, mmb::Document::kInPlaceDisabled);
         const bool validateForStorage = false;
         const FieldRefSet emptyImmutablePaths;
+        const bool isInsert = false;
         BSONObj logObj;
         status = driver.update(
-            StringData(), &document, validateForStorage, emptyImmutablePaths, &logObj);
+            StringData(), &document, validateForStorage, emptyImmutablePaths, isInsert, &logObj);
         if (!status.isOK())
             return status;
         BSONObj newObj = document.getObject().copy();
@@ -226,7 +227,9 @@ Status AuthzManagerExternalStateMock::updateOne(OperationContext* opCtx,
 
         const bool validateForStorage = false;
         const FieldRefSet emptyImmutablePaths;
-        status = driver.update(StringData(), &document, validateForStorage, emptyImmutablePaths);
+        const bool isInsert = false;
+        status = driver.update(
+            StringData(), &document, validateForStorage, emptyImmutablePaths, isInsert);
         if (!status.isOK()) {
             return status;
         }
diff --git a/src/mongo/db/auth/role_graph_update.cpp b/src/mongo/db/auth/role_graph_update.cpp
index 0515692cf48..12e7a098784 100644
--- a/src/mongo/db/auth/role_graph_update.cpp
+++ b/src/mongo/db/auth/role_graph_update.cpp
@@ -227,7 +227,9 @@ Status handleOplogUpdate(OperationContext* opCtx,
 
     const bool validateForStorage = false;
     const FieldRefSet emptyImmutablePaths;
-    status = driver.update(StringData(), &roleDocument, validateForStorage, emptyImmutablePaths);
+    bool isInsert = false;
+    status = driver.update(
+        StringData(), &roleDocument, validateForStorage, emptyImmutablePaths, isInsert);
     if (!status.isOK())
         return status;
 
diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index 83f07245294..0dfaacddc87 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -221,6 +221,7 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
 
     Status status = Status::OK();
     const bool validateForStorage = getOpCtx()->writesAreReplicated() && _enforceOkForStorage;
+    const bool isInsert = false;
     FieldRefSet immutablePaths;
     if (getOpCtx()->writesAreReplicated() && !request->isFromMigration()) {
         auto immutablePathsVector = getImmutableFields(getOpCtx(), request->getNamespaceString());
@@ -232,8 +233,13 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
     }
     if (!driver->needMatchDetails()) {
         // If we don't need match details, avoid doing the rematch
-        status = driver->update(
-            StringData(), &_doc, validateForStorage, immutablePaths, &logObj, &docWasModified);
+        status = driver->update(StringData(),
+                                &_doc,
+                                validateForStorage,
+                                immutablePaths,
+                                isInsert,
+                                &logObj,
+                                &docWasModified);
     } else {
         // If there was a matched field, obtain it.
         MatchDetails matchDetails;
@@ -246,8 +252,13 @@ BSONObj UpdateStage::transformAndUpdate(const Snapshotted<BSONObj>& oldObj, Reco
         if (matchDetails.hasElemMatchKey())
             matchedField = matchDetails.elemMatchKey();
 
-        status = driver->update(
-            matchedField, &_doc, validateForStorage, immutablePaths, &logObj, &docWasModified);
+        status = driver->update(matchedField,
+                                &_doc,
+                                validateForStorage,
+                                immutablePaths,
+                                isInsert,
+                                &logObj,
+                                &docWasModified);
     }
 
     if (!status.isOK()) {
@@ -379,7 +390,6 @@ BSONObj UpdateStage::applyUpdateOpsForInsert(OperationContext* opCtx,
     // oplog record, then. We also set the context of the update driver to the INSERT_CONTEXT.
     // Some mods may only work in that context (e.g. $setOnInsert).
     driver->setLogOp(false);
-    driver->setInsert(true);
 
     FieldRefSet immutablePaths;
     if (!isInternalRequest) {
@@ -404,10 +414,12 @@ BSONObj UpdateStage::applyUpdateOpsForInsert(OperationContext* opCtx,
     // Apply the update modifications here. Do not validate for storage, since we will validate the
     // entire document after the update. However, we ensure that no immutable fields are updated.
     const bool validateForStorage = false;
+    const bool isInsert = true;
     if (isInternalRequest) {
         immutablePaths.clear();
     }
-    Status updateStatus = driver->update(StringData(), doc, validateForStorage, immutablePaths);
+    Status updateStatus =
+        driver->update(StringData(), doc, validateForStorage, immutablePaths, isInsert);
     if (!updateStatus.isOK()) {
         uasserted(16836, updateStatus.reason());
     }
diff --git a/src/mongo/db/ops/update.cpp b/src/mongo/db/ops/update.cpp
index ff5ad3b307d..8e0eb55a9ec 100644
--- a/src/mongo/db/ops/update.cpp
+++ b/src/mongo/db/ops/update.cpp
@@ -121,8 +121,10 @@ BSONObj applyUpdateOperators(OperationContext* opCtx,
 
     const bool validateForStorage = false;
     const FieldRefSet emptyImmutablePaths;
+    const bool isInsert = false;
 
-    uassertStatusOK(driver.update(StringData(), &doc, validateForStorage, emptyImmutablePaths));
+    uassertStatusOK(
+        driver.update(StringData(), &doc, validateForStorage, emptyImmutablePaths, isInsert));
 
     return doc.getObject();
 }
diff --git a/src/mongo/db/update/update_driver.cpp b/src/mongo/db/update/update_driver.cpp
index 1b6e5cfcfd8..8eff56949f0 100644
--- a/src/mongo/db/update/update_driver.cpp
+++ b/src/mongo/db/update/update_driver.cpp
@@ -244,6 +244,7 @@ Status UpdateDriver::update(StringData matchedField,
                             mutablebson::Document* doc,
                             bool validateForStorage,
                             const FieldRefSet& immutablePaths,
+                            bool isInsert,
                             BSONObj* logOpRec,
                             bool* docWasModified,
                             FieldRefSetWithStorage* modifiedPaths) {
@@ -256,7 +257,7 @@ Status UpdateDriver::update(StringData matchedField,
 
     UpdateNode::ApplyParams applyParams(doc->root(), immutablePaths);
     applyParams.matchedField = matchedField;
-    applyParams.insert = _insert;
+    applyParams.insert = isInsert;
     applyParams.fromOplogApplication = _fromOplogApplication;
     applyParams.validateForStorage = validateForStorage;
     applyParams.indexData = _indexedFields;
diff --git a/src/mongo/db/update/update_driver.h b/src/mongo/db/update/update_driver.h
index bd2c94f7cc5..657b6fd3d86 100644
--- a/src/mongo/db/update/update_driver.h
+++ b/src/mongo/db/update/update_driver.h
@@ -96,6 +96,8 @@ public:
      * constraints. Ensures that no paths in 'immutablePaths' are modified (though they may be
      * created, if they do not yet exist).
      *
+     * The value of 'isInsert' controls whether $setOnInsert modifiers get applied.
+     *
      * If 'modifiedPaths' is not null, this method will populate it with the set of paths that were
      * either modified at runtime or present statically in the update modifiers. For arrays, the
      * set will include only the path to the array if the length has changed. All paths encode array
@@ -108,6 +110,7 @@ public:
                   mutablebson::Document* doc,
                   bool validateForStorage,
                   const FieldRefSet& immutablePaths,
+                  bool isInsert,
                   BSONObj* logOpRec = nullptr,
                   bool* docWasModified = nullptr,
                   FieldRefSetWithStorage* modifiedPaths = nullptr);
@@ -128,10 +131,6 @@ public:
     bool fromOplogApplication() const;
     void setFromOplogApplication(bool fromOplogApplication);
 
-    void setInsert(bool insert) {
-        _insert = insert;
-    }
-
     mutablebson::Document& getDocument() {
         return _objDoc;
     }
@@ -191,9 +190,6 @@ private:
     // Do any of the mods require positional match details when calling 'prepare'?
     bool _positional = false;
 
-    // Is this update going to be an upsert?
-    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_driver_test.cpp b/src/mongo/db/update/update_driver_test.cpp
index 44a44514d8b..b64aa84bfd9 100644
--- a/src/mongo/db/update/update_driver_test.cpp
+++ b/src/mongo/db/update/update_driver_test.cpp
@@ -156,11 +156,12 @@ TEST(Collator, SetCollationUpdatesModifierInterfaces) {
 
     const bool validateForStorage = true;
     const FieldRefSet emptyImmutablePaths;
+    const bool isInsert = false;
     bool modified = false;
     mutablebson::Document doc(fromjson("{a: 'cba'}"));
     driver.setCollator(&reverseStringCollator);
-    driver.update(StringData(), &doc, validateForStorage, emptyImmutablePaths, nullptr, &modified)
-        .transitional_ignore();
+    ASSERT_OK(driver.update(
+        StringData(), &doc, validateForStorage, emptyImmutablePaths, isInsert, nullptr, &modified));
 
     ASSERT_TRUE(modified);
 }
@@ -548,11 +549,13 @@ public:
 
         const bool validateForStorage = true;
         const FieldRefSet emptyImmutablePaths;
+        const bool isInsert = false;
         FieldRefSetWithStorage modifiedPaths;
         ASSERT_OK(driver.update(matchedField,
                                 doc,
                                 validateForStorage,
                                 emptyImmutablePaths,
+                                isInsert,
                                 nullptr,
                                 nullptr,
                                 &modifiedPaths));