SERVER-40185 Replacement upsert in Stitch Lib takes '_id' from query.

2 files changed, 16 insertions, 1 deletions

diff --git a/src/mongo/embedded/stitch_support/stitch_support.cpp b/src/mongo/embedded/stitch_support/stitch_support.cpp
index 14c046c7d01..06ba765eab8 100644
--- a/src/mongo/embedded/stitch_support/stitch_support.cpp
+++ b/src/mongo/embedded/stitch_support/stitch_support.cpp
@@ -619,12 +619,15 @@ 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);
 
+        const mongo::FieldRef idFieldRef("_id");
+        mongo::FieldRefSet immutablePaths;
+        invariant(immutablePaths.insert(&idFieldRef));
+
         if (update->matcher) {
             uassertStatusOK(update->updateDriver.populateDocumentWithQueryFields(
                 update->opCtx.get(),
diff --git a/src/mongo/embedded/stitch_support/stitch_support_test.cpp b/src/mongo/embedded/stitch_support/stitch_support_test.cpp
index c81ff5a8ced..e73ef6627e8 100644
--- a/src/mongo/embedded/stitch_support/stitch_support_test.cpp
+++ b/src/mongo/embedded/stitch_support/stitch_support_test.cpp
@@ -474,6 +474,10 @@ TEST_F(StitchSupportTest, TestReplacementStyleUpdateReportsNoModifiedPaths) {
     ASSERT_EQ("[]", getModifiedPaths());
 }
 
+TEST_F(StitchSupportTest, TestReplacementStyleUpdatePreservesId) {
+    ASSERT_EQ("{ \"_id\" : 123, \"b\" : 789 }", checkUpdate("{b: 789}", "{_id: 123, a: 456}"));
+}
+
 TEST_F(StitchSupportTest, TestUpdateArrayElement) {
     ASSERT_EQ("{ \"a\" : [ 2, 2 ] }", checkUpdate("{$set: {'a.0': 2}}", "{a: [1, 2]}"));
     ASSERT_EQ("[a.0]", getModifiedPaths());
@@ -580,6 +584,14 @@ TEST_F(StitchSupportTest, TestUpsertEmptyMatcher) {
     ASSERT_EQ("{ \"a\" : [ { \"b\" : 2 }, false ] }", checkUpsert("{a: [{b: 2}, false]}", "{}"));
 }
 
+TEST_F(StitchSupportTest, TestUpsertWithReplacementUpdate) {
+    ASSERT_EQ("{ \"_id\" : 1, \"a\" : 2 }", checkUpsert("{a: 2}", "{_id: 1}"));
+    ASSERT_EQ("{ \"_id\" : 1, \"a\" : 2 }", checkUpsert("{a: 2}", "{$and: [{_id: 1}]}"));
+
+    // Upsert with replacement update ues the '_id' field from the query but not any other fields.
+    ASSERT_EQ("{ \"_id\" : 1, \"b\" : 4 }", checkUpsert("{b: 4}", "{_id: 1, a: 2, b: 3}"));
+}
+
 TEST_F(StitchSupportTest, TestUpsertProducesProperStatus) {
     ASSERT_EQ("Cannot apply array updates to non-array element a: 1",
               checkUpsertStatus("{$set: {'a.$[].b': 1}}", "{a: 1}"));