fix(query): handle casting array filter paths underneath array filter paths with embedded discriminators

Fix #15386

Author: vkarpov15

lib/helpers/query/getEmbeddedDiscriminatorPath.js

@@ -17,7 +17,7 @@ const updatedPathsByArrayFilter = require('../update/updatedPathsByArrayFilter')
  */
 
 module.exports = function getEmbeddedDiscriminatorPath(schema, update, filter, path, options) {
-  const parts = path.split('.');
+  const parts = path.indexOf('.') === -1 ? [path] : path.split('.');
   let schematype = null;
   let type = 'adhocOrUndefined';
 
@@ -26,9 +26,10 @@ module.exports = function getEmbeddedDiscriminatorPath(schema, update, filter, p
   const arrayFilters = options != null && Array.isArray(options.arrayFilters) ?
     options.arrayFilters : [];
   const updatedPathsByFilter = updatedPathsByArrayFilter(update);
+  let startIndex = 0;
 
   for (let i = 0; i < parts.length; ++i) {
-    const originalSubpath = parts.slice(0, i + 1).join('.');
+    const originalSubpath = parts.slice(startIndex, i + 1).join('.');
     const subpath = cleanPositionalOperators(originalSubpath);
     schematype = schema.path(subpath);
     if (schematype == null) {
@@ -89,6 +90,8 @@ module.exports = function getEmbeddedDiscriminatorPath(schema, update, filter, p
 
       const rest = parts.slice(i + 1).join('.');
       schematype = discriminatorSchema.path(rest);
+      schema = discriminatorSchema;
+      startIndex = i + 1;
       if (schematype != null) {
         type = discriminatorSchema._getPathType(rest);
         break;

lib/helpers/schema/getPath.js

@@ -8,7 +8,7 @@ const numberRE = /^\d+$/;
  * @api private
  */
 
-module.exports = function getPath(schema, path) {
+module.exports = function getPath(schema, path, discriminatorValueMap) {
   let schematype = schema.path(path);
   if (schematype != null) {
     return schematype;
@@ -26,10 +26,14 @@ module.exports = function getPath(schema, path) {
     schematype = schema.path(cur);
     if (schematype != null && schematype.schema) {
       schema = schematype.schema;
+      const currentPath = cur;
       cur = '';
       if (!isArray && schematype.$isMongooseDocumentArray) {
         isArray = true;
       }
+      if (discriminatorValueMap && discriminatorValueMap[currentPath]) {
+        schema = schema.discriminators[discriminatorValueMap[currentPath]] ?? schema;
+      }
     }
   }
 

lib/helpers/update/castArrayFilters.js

@@ -33,6 +33,9 @@ function _castArrayFilters(arrayFilters, schema, strictQuery, updatedPathsByFilt
     return;
   }
 
+  // Map of embedded discriminator values set in the array filters
+  const discriminatorValueMap = {};
+
   for (const filter of arrayFilters) {
     if (filter == null) {
       throw new Error(`Got null array filter in ${arrayFilters}`);
@@ -58,12 +61,13 @@ function _castArrayFilters(arrayFilters, schema, strictQuery, updatedPathsByFilt
       updatedPathsByFilter[filterWildcardPath]
     );
 
-    const baseSchematype = getPath(schema, baseFilterPath);
+    const baseSchematype = getPath(schema, baseFilterPath, discriminatorValueMap);
     let filterBaseSchema = baseSchematype != null ? baseSchematype.schema : null;
     if (filterBaseSchema != null &&
         filterBaseSchema.discriminators != null &&
         filter[filterWildcardPath + '.' + filterBaseSchema.options.discriminatorKey]) {
       filterBaseSchema = filterBaseSchema.discriminators[filter[filterWildcardPath + '.' + filterBaseSchema.options.discriminatorKey]] || filterBaseSchema;
+      discriminatorValueMap[baseFilterPath] = filter[filterWildcardPath + '.' + filterBaseSchema.options.discriminatorKey];
     }
 
     for (const key of keys) {
@@ -83,7 +87,7 @@ function _castArrayFilters(arrayFilters, schema, strictQuery, updatedPathsByFilt
         // If there are multiple array filters in the path being updated, make sure
         // to replace them so we can get the schema path.
         filterPathRelativeToBase = cleanPositionalOperators(filterPathRelativeToBase);
-        schematype = getPath(filterBaseSchema, filterPathRelativeToBase);
+        schematype = getPath(filterBaseSchema, filterPathRelativeToBase, discriminatorValueMap);
       }
 
       if (schematype == null) {

test/helpers/update.castArrayFilters.test.js

@@ -349,4 +349,52 @@ describe('castArrayFilters', function() {
 
     assert.strictEqual(q.getUpdate().$set['groups.$[group].tags.$[tag]'], '42');
   });
+
+  it('casts paths underneath embedded discriminators (gh-15386)', async function() {
+    const eventSchema = new Schema({ message: String }, { discriminatorKey: 'kind', _id: false });
+    const batchSchema = new Schema({ events: [eventSchema] });
+
+    const docArray = batchSchema.path('events');
+    docArray.discriminator('Clicked', new Schema({ element: { type: String, required: true } }, { _id: false }));
+
+    const productSchema = new Schema({
+      name: String,
+      price: Number
+    });
+
+    docArray.discriminator(
+      'Purchased',
+      new Schema({
+        products: {
+          type: [productSchema],
+          required: true
+        }
+      })
+    );
+
+    const q = new Query();
+    q.schema = batchSchema;
+
+    const filter = {};
+    const update = {
+      $set: {
+        'events.$[event].products.$[product].price': '20'
+      }
+    };
+    const purchasedId = new Types.ObjectId();
+    const productId = new Types.ObjectId();
+    const opts = {
+      arrayFilters: [
+        { 'event._id': purchasedId, 'event.kind': 'Purchased' },
+        { 'product._id': productId.toString() }
+      ]
+    };
+
+    q.updateOne(filter, update, opts);
+    castArrayFilters(q);
+    q._update = q._castUpdate(q._update, false);
+
+    assert.strictEqual(q.getOptions().arrayFilters[1]['product._id'].toHexString(), productId.toHexString());
+    assert.strictEqual(q.getUpdate().$set['events.$[event].products.$[product].price'], 20);
+  });
 });