Merge pull request #13981 from Automattic/vkarpov15/gh-13944

fix(model): add versionKey to bulkWrite when inserting or upserting
Automattic · Oct 17, 2023 · 68c50c2 · 68c50c2
2 parents bbe3a0b + fa05928
commit 68c50c2
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 26 deletions.
diff --git a/lib/helpers/model/castBulkWrite.js b/lib/helpers/model/castBulkWrite.js
@@ -6,6 +6,7 @@ const applyTimestampsToChildren = require('../update/applyTimestampsToChildren')
 const applyTimestampsToUpdate = require('../update/applyTimestampsToUpdate');
 const cast = require('../../cast');
 const castUpdate = require('../query/castUpdate');
+const decorateUpdateWithVersionKey = require('../update/decorateUpdateWithVersionKey');
 const { inspect } = require('util');
 const setDefaultsOnInsert = require('../setDefaultsOnInsert');
 
@@ -33,6 +34,10 @@ module.exports = function castBulkWrite(originalModel, op, options) {
       if (options.session != null) {
         doc.$session(options.session);
       }
+      const versionKey = model?.schema?.options?.versionKey;
+      if (versionKey && doc[versionKey] == null) {
+        doc[versionKey] = 0;
+      }
       op['insertOne']['document'] = doc;
 
       if (options.skipValidation || op['insertOne'].skipValidation) {
@@ -81,6 +86,12 @@ module.exports = function castBulkWrite(originalModel, op, options) {
           });
         }
 
+        decorateUpdateWithVersionKey(
+          op['updateOne']['update'],
+          op['updateOne'],
+          model.schema.options.versionKey
+        );
+
         op['updateOne']['filter'] = cast(model.schema, op['updateOne']['filter'], {
           strict: strict,
           upsert: op['updateOne'].upsert
@@ -133,6 +144,12 @@ module.exports = function castBulkWrite(originalModel, op, options) {
 
         _addDiscriminatorToObject(schema, op['updateMany']['filter']);
 
+        decorateUpdateWithVersionKey(
+          op['updateMany']['update'],
+          op['updateMany'],
+          model.schema.options.versionKey
+        );
+
         op['updateMany']['filter'] = cast(model.schema, op['updateMany']['filter'], {
           strict: strict,
           upsert: op['updateMany'].upsert
@@ -173,6 +190,10 @@ module.exports = function castBulkWrite(originalModel, op, options) {
       if (options.session != null) {
         doc.$session(options.session);
       }
+      const versionKey = model?.schema?.options?.versionKey;
+      if (versionKey && doc[versionKey] == null) {
+        doc[versionKey] = 0;
+      }
       op['replaceOne']['replacement'] = doc;
 
       if (options.skipValidation || op['replaceOne'].skipValidation) {

diff --git a/lib/helpers/update/decorateUpdateWithVersionKey.js b/lib/helpers/update/decorateUpdateWithVersionKey.js
@@ -0,0 +1,26 @@
+'use strict';
+
+const modifiedPaths = require('./modifiedPaths');
+
+/**
+ * Decorate the update with a version key, if necessary
+ * @api private
+ */
+
+module.exports = function decorateUpdateWithVersionKey(update, options, versionKey) {
+  if (!versionKey || !(options && options.upsert || false)) {
+    return;
+  }
+
+  const updatedPaths = modifiedPaths(update);
+  if (!updatedPaths[versionKey]) {
+    if (options.overwrite) {
+      update[versionKey] = 0;
+    } else {
+      if (!update.$setOnInsert) {
+        update.$setOnInsert = {};
+      }
+      update.$setOnInsert[versionKey] = 0;
+    }
+  }
+};
diff --git a/lib/model.js b/lib/model.js
@@ -33,6 +33,7 @@ const assignVals = require('./helpers/populate/assignVals');
 const castBulkWrite = require('./helpers/model/castBulkWrite');
 const clone = require('./helpers/clone');
 const createPopulateQueryFilter = require('./helpers/populate/createPopulateQueryFilter');
+const decorateUpdateWithVersionKey = require('./helpers/update/decorateUpdateWithVersionKey');
 const getDefaultBulkwriteResult = require('./helpers/getDefaultBulkwriteResult');
 const getSchemaDiscriminatorByValue = require('./helpers/discriminator/getSchemaDiscriminatorByValue');
 const discriminator = require('./helpers/model/discriminator');
@@ -54,7 +55,6 @@ const isPathExcluded = require('./helpers/projection/isPathExcluded');
 const decorateDiscriminatorIndexOptions = require('./helpers/indexes/decorateDiscriminatorIndexOptions');
 const isPathSelectedInclusive = require('./helpers/projection/isPathSelectedInclusive');
 const leanPopulateMap = require('./helpers/populate/leanPopulateMap');
-const modifiedPaths = require('./helpers/update/modifiedPaths');
 const parallelLimit = require('./helpers/parallelLimit');
 const parentPaths = require('./helpers/path/parentPaths');
 const prepareDiscriminatorPipeline = require('./helpers/aggregate/prepareDiscriminatorPipeline');
@@ -2451,37 +2451,14 @@ Model.findOneAndUpdate = function(conditions, update, options) {
     _isNested: true
   });
 
-  _decorateUpdateWithVersionKey(update, options, this.schema.options.versionKey);
+  decorateUpdateWithVersionKey(update, options, this.schema.options.versionKey);
 
   const mq = new this.Query({}, {}, this, this.$__collection);
   mq.select(fields);
 
   return mq.findOneAndUpdate(conditions, update, options);
 };
 
-/**
- * Decorate the update with a version key, if necessary
- * @api private
- */
-
-function _decorateUpdateWithVersionKey(update, options, versionKey) {
-  if (!versionKey || !(options && options.upsert || false)) {
-    return;
-  }
-
-  const updatedPaths = modifiedPaths(update);
-  if (!updatedPaths[versionKey]) {
-    if (options.overwrite) {
-      update[versionKey] = 0;
-    } else {
-      if (!update.$setOnInsert) {
-        update.$setOnInsert = {};
-      }
-      update.$setOnInsert[versionKey] = 0;
-    }
-  }
-}
-
 /**
  * Issues a mongodb findOneAndUpdate command by a document's _id field.
  * `findByIdAndUpdate(id, ...)` is equivalent to `findOneAndUpdate({ _id: id }, ...)`.
@@ -4022,7 +3999,7 @@ function _update(model, op, conditions, doc, options) {
   model.schema &&
   model.schema.options &&
   model.schema.options.versionKey || null;
-  _decorateUpdateWithVersionKey(doc, options, versionKey);
+  decorateUpdateWithVersionKey(doc, options, versionKey);
 
   return mq[op](conditions, doc, options);
 }

diff --git a/test/model.test.js b/test/model.test.js
@@ -4043,6 +4043,49 @@ describe('Model', function() {
 
       });
 
+      it('sets version key (gh-13944)', async function() {
+        const userSchema = new Schema({
+          firstName: { type: String, required: true },
+          lastName: { type: String }
+        });
+        const User = db.model('User', userSchema);
+
+        await User.bulkWrite([
+          {
+            updateOne: {
+              filter: { lastName: 'Gibbons' },
+              update: { firstName: 'Peter' },
+              upsert: true
+            }
+          },
+          {
+            insertOne: {
+              document: {
+                firstName: 'Michael',
+                lastName: 'Bolton'
+              }
+            }
+          },
+          {
+            replaceOne: {
+              filter: { lastName: 'Lumbergh' },
+              replacement: { firstName: 'Bill', lastName: 'Lumbergh' },
+              upsert: true
+            }
+          }
+        ], { ordered: false });
+
+        const users = await User.find();
+        assert.deepStrictEqual(
+          users.map(user => user.firstName).sort(),
+          ['Bill', 'Michael', 'Peter']
+        );
+        assert.deepStrictEqual(
+          users.map(user => user.__v),
+          [0, 0, 0]
+        );
+      });
+
       it('with single nested and setOnInsert (gh-7534)', function() {
         const nested = new Schema({ name: String });
         const schema = new Schema({ nested: nested });