Merge pull request #14687 from Automattic/vkarpov15/gh-14678

fix: handle casting primitive array with $elemMatch in bulkWrite()
Automattic · Jun 25, 2024 · 31111f2 · 31111f2
2 parents 516464a + 091c85d
commit 31111f2
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 19 deletions.
diff --git a/lib/schema/array.js b/lib/schema/array.js
@@ -607,24 +607,7 @@ function cast$elemMatch(val, context) {
     }
   }
 
-  // Is this an embedded discriminator and is the discriminator key set?
-  // If so, use the discriminator schema. See gh-7449
-  const discriminatorKey = this &&
-    this.casterConstructor &&
-    this.casterConstructor.schema &&
-    this.casterConstructor.schema.options &&
-    this.casterConstructor.schema.options.discriminatorKey;
-  const discriminators = this &&
-  this.casterConstructor &&
-  this.casterConstructor.schema &&
-  this.casterConstructor.schema.discriminators || {};
-  if (discriminatorKey != null &&
-      val[discriminatorKey] != null &&
-      discriminators[val[discriminatorKey]] != null) {
-    return cast(discriminators[val[discriminatorKey]], val, null, this && this.$$context);
-  }
-  const schema = this.casterConstructor.schema ?? context.schema;
-  return cast(schema, val, null, this && this.$$context);
+  return val;
 }
 
 const handle = SchemaArray.prototype.$conditionalHandlers = {};

diff --git a/lib/schema/documentArray.js b/lib/schema/documentArray.js
@@ -11,9 +11,11 @@ const SchemaArray = require('./array');
 const SchemaDocumentArrayOptions =
   require('../options/schemaDocumentArrayOptions');
 const SchemaType = require('../schemaType');
+const cast = require('../cast');
 const discriminator = require('../helpers/model/discriminator');
 const handleIdOption = require('../helpers/schema/handleIdOption');
 const handleSpreadDoc = require('../helpers/document/handleSpreadDoc');
+const isOperator = require('../helpers/query/isOperator');
 const utils = require('../utils');
 const getConstructor = require('../helpers/discriminator/getConstructor');
 const InvalidSchemaOptionError = require('../error/invalidSchemaOption');
@@ -114,6 +116,7 @@ SchemaDocumentArray.options = { castNonArrays: true };
 SchemaDocumentArray.prototype = Object.create(SchemaArray.prototype);
 SchemaDocumentArray.prototype.constructor = SchemaDocumentArray;
 SchemaDocumentArray.prototype.OptionsConstructor = SchemaDocumentArrayOptions;
+SchemaDocumentArray.prototype.$conditionalHandlers = { ...SchemaArray.prototype.$conditionalHandlers };
 
 /*!
  * ignore
@@ -609,6 +612,44 @@ SchemaDocumentArray.setters = [];
 
 SchemaDocumentArray.get = SchemaType.get;
 
+/*!
+ * Handle casting $elemMatch operators
+ */
+
+SchemaDocumentArray.prototype.$conditionalHandlers.$elemMatch = cast$elemMatch;
+
+function cast$elemMatch(val, context) {
+  const keys = Object.keys(val);
+  const numKeys = keys.length;
+  for (let i = 0; i < numKeys; ++i) {
+    const key = keys[i];
+    const value = val[key];
+    if (isOperator(key) && value != null) {
+      val[key] = this.castForQuery(key, value, context);
+    }
+  }
+
+  // Is this an embedded discriminator and is the discriminator key set?
+  // If so, use the discriminator schema. See gh-7449
+  const discriminatorKey = this &&
+    this.casterConstructor &&
+    this.casterConstructor.schema &&
+    this.casterConstructor.schema.options &&
+    this.casterConstructor.schema.options.discriminatorKey;
+  const discriminators = this &&
+  this.casterConstructor &&
+  this.casterConstructor.schema &&
+  this.casterConstructor.schema.discriminators || {};
+  if (discriminatorKey != null &&
+      val[discriminatorKey] != null &&
+      discriminators[val[discriminatorKey]] != null) {
+    return cast(discriminators[val[discriminatorKey]], val, null, this && this.$$context);
+  }
+
+  const schema = this.casterConstructor.schema ?? context.schema;
+  return cast(schema, val, null, this && this.$$context);
+}
+
 /*!
  * Module exports.
  */

diff --git a/test/model.query.casting.test.js b/test/model.query.casting.test.js
@@ -453,7 +453,10 @@ describe('model query casting', function() {
       const id = post._id.toString();
 
       await post.save();
-      const doc = await BlogPostB.findOne({ _id: id, comments: { $not: { $elemMatch: { _id: commentId.toString() } } } });
+      const doc = await BlogPostB.findOne({
+        _id: id,
+        comments: { $not: { $elemMatch: { _id: commentId.toString() } } }
+      });
       assert.equal(doc, null);
     });
 

diff --git a/test/model.test.js b/test/model.test.js
@@ -4470,6 +4470,43 @@ describe('Model', function() {
         assert.equal(err.validationErrors[0].path, 'age');
         assert.equal(err.results[0].path, 'age');
       });
+
+      it('casts $elemMatch filter (gh-14678)', async function() {
+        const schema = new mongoose.Schema({
+          name: String,
+          ids: [String]
+        });
+        const TestModel = db.model('Test', schema);
+
+        const { _id } = await TestModel.create({ ids: ['1'] });
+        await TestModel.bulkWrite([
+          {
+            updateOne: {
+              filter: {
+                ids: {
+                  $elemMatch: {
+                    $in: [1]
+                  }
+                }
+              },
+              update: {
+                $set: {
+                  name: 'test'
+                },
+                $addToSet: {
+                  ids: {
+                    $each: [1, '2', 3]
+                  }
+                }
+              }
+            }
+          }
+        ]);
+
+        const doc = await TestModel.findById(_id).orFail();
+        assert.strictEqual(doc.name, 'test');
+        assert.deepStrictEqual(doc.ids, ['1', '2', '3']);
+      });
     });
 
     it('deleteOne with cast error (gh-5323)', async function() {