feat(connection): make transaction() reset versionKey if transactio…

…n fails Re: #8380
Automattic · Jun 15, 2020 · 688d19d · 688d19d
1 parent 78f5dbb
commit 688d19d
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 8 deletions.
diff --git a/lib/connection.js b/lib/connection.js
@@ -452,7 +452,7 @@ Connection.prototype.startSession = _wrapConnHelper(function startSession(option
 
 Connection.prototype.transaction = function transaction(fn) {
   return this.startSession().then(session => {
-    session[sessionNewDocuments] = [];
+    session[sessionNewDocuments] = new Map();
     return session.withTransaction(() => fn(session)).
       then(res => {
         delete session[sessionNewDocuments];
@@ -461,8 +461,13 @@ Connection.prototype.transaction = function transaction(fn) {
       catch(err => {
         // If transaction was aborted, we need to reset newly
         // inserted documents' `isNew`.
-        for (const doc of session[sessionNewDocuments]) {
-          doc.isNew = true;
+        for (const [doc, state] of session[sessionNewDocuments].entries()) {
+          if (state.hasOwnProperty('isNew')) {
+            doc.isNew = state.isNew;
+          }
+          if (state.hasOwnProperty('versionKey')) {
+            doc.set(doc.schema.options.versionKey, state.versionKey);
+          }
         }
         delete session[sessionNewDocuments];
         throw err;

diff --git a/lib/plugins/trackTransaction.js b/lib/plugins/trackTransaction.js
@@ -4,17 +4,24 @@ const sessionNewDocuments = require('../helpers/symbols').sessionNewDocuments;
 
 module.exports = function trackTransaction(schema) {
   schema.pre('save', function() {
-    if (!this.isNew) {
-      return;
-    }
-
     const session = this.$session();
     if (session == null) {
       return;
     }
     if (session.transaction == null || session[sessionNewDocuments] == null) {
       return;
     }
-    session[sessionNewDocuments].push(this);
+
+    if (!session[sessionNewDocuments].has(this)) {
+      const initialState = {};
+      if (this.isNew) {
+        initialState.isNew = true;
+      }
+      if (this.schema.options.versionKey) {
+        initialState.versionKey = this.get(this.schema.options.versionKey);
+      }
+
+      session[sessionNewDocuments].set(this, initialState);
+    }
   });
 };
diff --git a/test/docs/transactions.test.js b/test/docs/transactions.test.js
@@ -10,6 +10,7 @@ const Schema = mongoose.Schema;
 describe('transactions', function() {
   let db;
   let _skipped = false;
+  this.timeout(10000);
 
   before(function() {
     if (!process.env.REPLICA_SET) {
@@ -382,4 +383,25 @@ describe('transactions', function() {
       assert.ok(doc.isNew);
     });
   });
+
+  it('can save document after aborted transaction (gh-8380)', function() {
+    return co(function*() {
+      const schema = Schema({ name: String, arr: [String] });
+
+      const Test = db.model('gh8380', schema);
+
+      yield Test.createCollection();
+      yield Test.create({ name: 'foo', arr: ['bar'] });
+      const doc = yield Test.findOne();
+      yield db.
+        transaction(session => co(function*() {
+          doc.arr.pop();
+          yield doc.save({ session });
+          throw new Error('Oops');
+        })).
+        catch(err => assert.equal(err.message, 'Oops'));
+      doc.set('arr.0', 'qux');
+      yield doc.save();
+    });
+  });
 });