From 404b74bd0900b8ecd7a76832215f32851e6acf3f Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Fri, 7 Feb 2025 16:09:52 -0500
Subject: [PATCH] fix: infer discriminator key if set in `$set` with
 overwriteDiscriminatorKey

Re: #15218
---
 CHANGELOG.md                    |  2 --
 lib/helpers/query/castUpdate.js | 10 ++++++++++
 test/model.test.js              | 12 +++++++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5e986e8fad2..073c270e17a 100644
--- 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
index 3cf30cb0e17..f0819b3a586 100644
--- 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
index dd4d18aa1bf..96917d0e5ea 100644
--- 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);
       });