fix: cast using overwritten embedded discriminator key when set

Fix Automattic#15051

Author: vkarpov15

lib/helpers/query/getEmbeddedDiscriminatorPath.js

@@ -28,7 +28,8 @@ module.exports = function getEmbeddedDiscriminatorPath(schema, update, filter, p
   const updatedPathsByFilter = updatedPathsByArrayFilter(update);
 
   for (let i = 0; i < parts.length; ++i) {
-    const subpath = cleanPositionalOperators(parts.slice(0, i + 1).join('.'));
+    const originalSubpath = parts.slice(0, i + 1).join('.');
+    const subpath = cleanPositionalOperators(originalSubpath);
     schematype = schema.path(subpath);
     if (schematype == null) {
       continue;
@@ -56,6 +57,11 @@ module.exports = function getEmbeddedDiscriminatorPath(schema, update, filter, p
         discriminatorKey = filter[wrapperPath].$elemMatch[key];
       }
 
+      const discriminatorKeyUpdatePath = originalSubpath + '.' + key;
+      if (discriminatorKeyUpdatePath in update) {
+        discriminatorKey = update[discriminatorKeyUpdatePath];
+      }
+
       if (discriminatorValuePath in update) {
         discriminatorKey = update[discriminatorValuePath];
       }

test/model.updateOne.test.js

@@ -3091,6 +3091,57 @@ describe('model: updateOne: ', function() {
     assert.equal(doc.login.keys.length, 1);
     assert.equal(doc.login.keys[0].id, 'test2');
   });
+  it('casts using overwritten discriminator key schema (gh-15051)', async function() {
+    const embedDiscriminatorSchema = new mongoose.Schema({
+      field1: String
+    });
+    const embedDiscriminatorSchema2 = new mongoose.Schema({
+      field2: String
+    });
+    const embedSchema = new mongoose.Schema({
+      field: String,
+      key: String
+    }, { discriminatorKey: 'key' });
+    embedSchema.discriminator('Type1', embedDiscriminatorSchema);
+    embedSchema.discriminator('Type2', embedDiscriminatorSchema2);
+
+    const testSchema = new mongoose.Schema({
+      testArray: [embedSchema]
+    });
+
+    const TestModel = db.model('Test', testSchema);
+    const test = new TestModel({
+      testArray: [{
+        key: 'Type1',
+        field: 'field',
+        field1: 'field1'
+      }]
+    });
+    await test.save();
+
+    const field2update = 'field2 update';
+    await TestModel.updateOne(
+      { _id: test._id },
+      {
+        $set: {
+          'testArray.$[element].key': 'Type2',
+          'testArray.$[element].field2': field2update
+        }
+      },
+      {
+        arrayFilters: [
+          {
+            'element._id': test.testArray[0]._id
+          }
+        ],
+        overwriteDiscriminatorKey: true
+      }
+    );
+
+    const r2 = await TestModel.findById(test._id);
+    assert.equal(r2.testArray[0].key, 'Type2');
+    assert.equal(r2.testArray[0].field2, field2update);
+  });
 });
 
 async function delay(ms) {