SERVER-39060 Add upsert function to Stitch Support Library

10 files changed, 150 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));
diff --git a/src/mongo/embedded/stitch_support/stitch_support.cpp b/src/mongo/embedded/stitch_support/stitch_support.cpp
index 15ee4dcbb6f..c73a8ee7acd 100644
--- a/src/mongo/embedded/stitch_support/stitch_support.cpp
+++ b/src/mongo/embedded/stitch_support/stitch_support.cpp
@@ -570,6 +570,7 @@ stitch_support_v1_update_apply(stitch_support_v1_update* const update,
                                                     &mutableDoc,
                                                     false /* validateForStorage */,
                                                     immutablePaths,
+                                                    false /* isInsert */,
                                                     nullptr /* logOpRec*/,
                                                     &docWasModified,
                                                     &modifiedPaths));
@@ -591,6 +592,39 @@ stitch_support_v1_update_apply(stitch_support_v1_update* const update,
     });
 }
 
+uint8_t* MONGO_API_CALL stitch_support_v1_update_upsert(stitch_support_v1_update* const update,
+                                                        stitch_support_v1_status* status) {
+    return enterCXX(mongo::getStatusImpl(status), [&] {
+        mongo::FieldRefSet immutablePaths;  //  Empty set
+        bool docWasModified = false;
+
+        mongo::mutablebson::Document mutableDoc(mongo::BSONObj(),
+                                                mongo::mutablebson::Document::kInPlaceDisabled);
+
+        uassertStatusOK(update->updateDriver.populateDocumentWithQueryFields(
+            update->opCtx.get(), *update->matcher->matcher.getQuery(), immutablePaths, mutableDoc));
+
+        uassertStatusOK(update->updateDriver.update(mongo::StringData() /* matchedField */,
+                                                    &mutableDoc,
+                                                    false /* validateForStorage */,
+                                                    immutablePaths,
+                                                    true /* isInsert */,
+                                                    nullptr /* logOpRec */,
+                                                    &docWasModified,
+                                                    nullptr /*modifiedPaths*/));
+
+        auto outputObj = mutableDoc.getObject();
+        size_t outputSize = static_cast<size_t>(outputObj.objsize());
+        auto output = new (std::nothrow) char[outputSize];
+
+        uassert(
+            mongo::ErrorCodes::ExceededMemoryLimit, "Failed to allocate memory for upsert", output);
+
+        static_cast<void>(std::copy_n(outputObj.objdata(), outputSize, output));
+        return mongo::toInterfaceType(output);
+    });
+}
+
 stitch_support_v1_update_details* MONGO_API_CALL stitch_support_v1_update_details_create(void) {
     return new stitch_support_v1_update_details;
 };
diff --git a/src/mongo/embedded/stitch_support/stitch_support.h b/src/mongo/embedded/stitch_support/stitch_support.h
index 7e8b2969370..18f772d665d 100644
--- a/src/mongo/embedded/stitch_support/stitch_support.h
+++ b/src/mongo/embedded/stitch_support/stitch_support.h
@@ -509,6 +509,23 @@ stitch_support_v1_update_apply(stitch_support_v1_update* const update,
                                stitch_support_v1_status* status);
 
 /**
+ * Return the document that would result from an {upsert: true} update operation that must perform
+ * an upsert. The resulting document is based on the match predicate in the associated matcher and
+ * on the update document. There is no "input document," like in stitch_support_v1_apply(), because
+ * upserts occur when no input document is found by the match predicate.
+ *
+ * This operation requires an update object that has an associated matcher. Returns a pointer to the
+ * output buffer on success or, on error, returns NULL and populates the 'status' object. The caller
+ * is responsible for destroying the result buffer with stitch_support_v1_bson_free().
+ *
+ * Note that, although a database upsert operations guarantees that the inserted document will have
+ * an '_id' field, this function does not populate an '_id' field if one is not in the match or
+ * update document.
+ */
+STITCH_SUPPORT_API uint8_t* MONGO_API_CALL stitch_support_v1_update_upsert(
+    stitch_support_v1_update* const update, stitch_support_v1_status* status);
+
+/**
  * Free the memory of a BSON buffer returned by stitch_support_v1_projection_apply() or
  * stitch_support_v1_update_apply(). This function can be safely called on a NULL pointer.
  *
diff --git a/src/mongo/embedded/stitch_support/stitch_support_test.cpp b/src/mongo/embedded/stitch_support/stitch_support_test.cpp
index cdeed2c7f5b..f18ce3be7a8 100644
--- a/src/mongo/embedded/stitch_support/stitch_support_test.cpp
+++ b/src/mongo/embedded/stitch_support/stitch_support_test.cpp
@@ -296,6 +296,49 @@ protected:
         return ss.str();
     }
 
+    auto checkUpsert(const char* expr, const char* match) {
+        stitch_support_v1_matcher* matcher =
+            stitch_support_v1_matcher_create(lib, toBSONForAPI(match).first, nullptr, nullptr);
+        ASSERT(matcher);
+        ON_BLOCK_EXIT([matcher] { stitch_support_v1_matcher_destroy(matcher); });
+
+        stitch_support_v1_update* update = stitch_support_v1_update_create(
+            lib, toBSONForAPI(expr).first, nullptr, matcher, nullptr, status);
+        ON_BLOCK_EXIT([update] { stitch_support_v1_update_destroy(update); });
+
+        uint8_t* upsertResult = stitch_support_v1_update_upsert(update, status);
+        ASSERT(upsertResult);
+        ON_BLOCK_EXIT([upsertResult] { stitch_support_v1_bson_free(upsertResult); });
+
+        return std::string(fromBSONForAPI(upsertResult));
+    }
+
+    auto checkUpsertStatus(const char* expr, const char* match) {
+        auto updateStatus = stitch_support_v1_status_create();
+        ON_BLOCK_EXIT([updateStatus] { stitch_support_v1_status_destroy(updateStatus); });
+
+        stitch_support_v1_matcher* matcher =
+            stitch_support_v1_matcher_create(lib, toBSONForAPI(match).first, nullptr, nullptr);
+        ASSERT(matcher);
+        ON_BLOCK_EXIT([matcher] { stitch_support_v1_matcher_destroy(matcher); });
+
+        stitch_support_v1_update* update = stitch_support_v1_update_create(
+            lib, toBSONForAPI(expr).first, nullptr, matcher, nullptr, updateStatus);
+
+        if (!update) {
+            ASSERT_EQ(STITCH_SUPPORT_V1_ERROR_EXCEPTION,
+                      stitch_support_v1_status_get_error(updateStatus));
+            // Make sure that we get a proper code back but don't worry about its exact value.
+            ASSERT_NE(0, stitch_support_v1_status_get_code(updateStatus));
+        } else {
+            ON_BLOCK_EXIT([update] { stitch_support_v1_update_destroy(update); });
+            auto upsertResult = stitch_support_v1_update_upsert(update, updateStatus);
+            ASSERT_NE(0, stitch_support_v1_status_get_code(updateStatus));
+            ASSERT(!upsertResult);
+        }
+        return std::string(stitch_support_v1_status_get_explanation(updateStatus));
+    }
+
     stitch_support_v1_status* status = nullptr;
     stitch_support_v1_lib* lib = nullptr;
     stitch_support_v1_update_details* updateDetails = nullptr;
@@ -472,6 +515,11 @@ TEST_F(StitchSupportTest, TestUpdateRespectsTheCollation) {
     ASSERT_EQ("[a]", getModifiedPaths());
 }
 
+TEST_F(StitchSupportTest, TestUpdateWithSetOnInsert) {
+    ASSERT_EQ("{ \"a\" : 1 }", checkUpdate("{$setOnInsert: {a: 2}}", "{a: 1}"));
+    ASSERT_EQ("[a]", getModifiedPaths());
+}
+
 TEST_F(StitchSupportTest, TestUpdateProducesProperStatus) {
     ASSERT_EQ("Unknown modifier: $bogus", checkUpdateStatus("{$bogus: {a: 2}}", "{a: 1}"));
     ASSERT_EQ("Updating the path 'a' would create a conflict at 'a'",
@@ -485,6 +533,18 @@ TEST_F(StitchSupportTest, TestUpdateProducesProperStatus) {
                   "{$set: {'a.$[i]': 2, 'a.$[j]': 3}}", "{a: [0]}", nullptr, " [{i: 0}, {j:0}]"));
 }
 
+TEST_F(StitchSupportTest, TestUpsert) {
+    ASSERT_EQ("{ \"_id\" : 1, \"a\" : 1 }", checkUpsert("{$set: {a: 1}}", "{_id: 1}"));
+    ASSERT_EQ("{ \"a\" : 2 }", checkUpsert("{$set: {a: 2}}", "{a: 1}"));
+    ASSERT_EQ("{ \"_id\" : 1, \"a\" : 1 }", checkUpsert("{$setOnInsert: {a: 1}}", "{_id: 1}"));
+    ASSERT_EQ("{ \"_id\" : 1, \"b\" : 1 }", checkUpsert("{$inc: {b: 1}}", "{_id: 1, a: {$gt: 2}}"));
+}
+
+TEST_F(StitchSupportTest, TestUpsertProducesProperStatus) {
+    ASSERT_EQ("Cannot apply array updates to non-array element a: 1",
+              checkUpsertStatus("{$set: {'a.$[].b': 1}}", "{a: 1}"));
+}
+
 }  // namespace
 
 // Define main function as an entry to these tests.