Merge pull request #15243 from Automattic/vkarpov15/gh-15218

fix: infer discriminator key if set in `$set` with overwriteDiscriminatorKey
Automattic · Feb 8, 2025 · 1a2faa7 · 1a2faa7
2 parents 6838d48 + 404b74b
commit 1a2faa7
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -123,8 +123,6 @@
  * fix: disallow using $where in match
  * perf: cache results from getAllSubdocs() on saveOptions, only loop through known subdoc properties #15055 #15029
  * fix(model+query): support overwriteDiscriminatorKey for bulkWrite updateOne and updateMany, allow inferring discriminator key from update #15046 #15040
-=======
->>>>>>> 7.x
 
 7.8.3 / 2024-11-26
 ==================

diff --git a/lib/helpers/query/castUpdate.js b/lib/helpers/query/castUpdate.js
@@ -82,6 +82,16 @@ module.exports = function castUpdate(schema, obj, options, context, filter) {
     schema = schema.discriminators[discriminatorValue] ||
       (byValue && byValue.schema) ||
       schema;
+  } else if (schema != null &&
+      options.overwriteDiscriminatorKey &&
+      obj.$set != null &&
+      utils.hasUserDefinedProperty(obj.$set, schema.options.discriminatorKey) &&
+      schema.discriminators != null) {
+    const discriminatorValue = obj.$set[schema.options.discriminatorKey];
+    const byValue = getDiscriminatorByValue(context.model.discriminators, discriminatorValue);
+    schema = schema.discriminators[discriminatorValue] ||
+      (byValue && byValue.schema) ||
+      schema;
   }
 
   if (options.upsert) {

diff --git a/test/model.test.js b/test/model.test.js
@@ -4174,7 +4174,7 @@ describe('Model', function() {
         assert.strictEqual(r2.testArray[0].nonexistentProp, undefined);
       });
 
-      it('handles overwriteDiscriminatorKey (gh-15040)', async function() {
+      it('handles overwriteDiscriminatorKey (gh-15218) (gh-15040)', async function() {
         const dSchema1 = new mongoose.Schema({
           field1: String
         });
@@ -4202,7 +4202,7 @@ describe('Model', function() {
         assert.equal(r1.field1, 'field1');
         assert.equal(r1.key, type1Key);
 
-        const field2 = 'field2';
+        let field2 = 'field2';
         await TestModel.bulkWrite([{
           updateOne: {
             filter: { _id: r1._id },
@@ -4214,7 +4214,13 @@ describe('Model', function() {
           }
         }]);
 
-        const r2 = await TestModel.findById(r1._id);
+        let r2 = await TestModel.findById(r1._id);
+        assert.equal(r2.key, type2Key);
+        assert.equal(r2.field2, field2);
+
+        field2 = 'field2 updated again';
+        await TestModel.updateOne({ _id: r1._id }, { $set: { key: type2Key, field2 } }, { overwriteDiscriminatorKey: true });
+        r2 = await TestModel.findById(r1._id);
         assert.equal(r2.key, type2Key);
         assert.equal(r2.field2, field2);
       });