diff --git a/lib/connection.js b/lib/connection.js index 12f0517aa06..85a00c27bcb 100644 --- a/lib/connection.js +++ b/lib/connection.js @@ -22,6 +22,8 @@ const utils = require('./utils'); const parseConnectionString = require('mongodb/lib/core').parseConnectionString; +const sessionNewDocuments = require('./helpers/symbols').sessionNewDocuments; + let id = 0; /*! @@ -417,6 +419,57 @@ Connection.prototype.startSession = _wrapConnHelper(function startSession(option cb(null, session); }); +/** + * _Requires MongoDB >= 3.6.0._ Executes the wrapped async function + * in a transaction. Mongoose will commit the transaction if the + * async function executes successfully and attempt to retry if + * there was a retriable error. + * + * Calls the MongoDB driver's [`session.withTransaction()`](http://mongodb.github.io/node-mongodb-native/3.5/api/ClientSession.html#withTransaction), + * but also handles resetting Mongoose document state as shown below. + * + * ####Example: + * + * const doc = new Person({ name: 'Will Riker' }); + * await db.transaction(async function setRank(session) { + * doc.rank = 'Captain'; + * await doc.save({ session }); + * doc.isNew; // false + * + * // Throw an error to abort the transaction + * throw new Error('Oops!'); + * }).catch(() => {}); + * + * // true, `transaction()` reset the document's state because the + * // transaction was aborted. + * doc.isNew; + * + * @method transaction + * @param {Function} fn Function to execute in a transaction + * @return {Promise} promise that resolves to the returned value of `fn` + * @api public + */ + +Connection.prototype.transaction = function transaction(fn) { + return this.startSession().then(session => { + session[sessionNewDocuments] = []; + return session.withTransaction(() => fn(session)). + then(res => { + delete session[sessionNewDocuments]; + return res; + }). + catch(err => { + // If transaction was aborted, we need to reset newly + // inserted documents' `isNew`. + for (const doc of session[sessionNewDocuments]) { + doc.isNew = true; + } + delete session[sessionNewDocuments]; + throw err; + }); + }); +}; + /** * Helper for `dropCollection()`. Will delete the given collection, including * all documents and indexes. diff --git a/lib/helpers/symbols.js b/lib/helpers/symbols.js index 3860f237743..4db388edfe7 100644 --- a/lib/helpers/symbols.js +++ b/lib/helpers/symbols.js @@ -11,4 +11,5 @@ exports.modelSymbol = Symbol('mongoose#Model'); exports.objectIdSymbol = Symbol('mongoose#ObjectId'); exports.populateModelSymbol = Symbol('mongoose.PopulateOptions#Model'); exports.schemaTypeSymbol = Symbol('mongoose#schemaType'); +exports.sessionNewDocuments = Symbol('mongoose:ClientSession#newDocuments'); exports.validatorErrorSymbol = Symbol('mongoose:validatorError'); \ No newline at end of file diff --git a/lib/index.js b/lib/index.js index 8c7c9f0aa90..b0d3e086151 100644 --- a/lib/index.js +++ b/lib/index.js @@ -34,6 +34,7 @@ const pkg = require('../package.json'); const cast = require('./cast'); const removeSubdocs = require('./plugins/removeSubdocs'); const saveSubdocs = require('./plugins/saveSubdocs'); +const trackTransaction = require('./plugins/trackTransaction'); const validateBeforeSave = require('./plugins/validateBeforeSave'); const Aggregate = require('./aggregate'); @@ -106,7 +107,8 @@ function Mongoose(options) { [saveSubdocs, { deduplicate: true }], [validateBeforeSave, { deduplicate: true }], [shardingPlugin, { deduplicate: true }], - [removeSubdocs, { deduplicate: true }] + [removeSubdocs, { deduplicate: true }], + [trackTransaction, { deduplicate: true }] ] }); } diff --git a/lib/plugins/trackTransaction.js b/lib/plugins/trackTransaction.js new file mode 100644 index 00000000000..c307f9b759a --- /dev/null +++ b/lib/plugins/trackTransaction.js @@ -0,0 +1,20 @@ +'use strict'; + +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); + }); +}; \ No newline at end of file diff --git a/test/docs/transactions.test.js b/test/docs/transactions.test.js index e27662861d8..73daab51e83 100644 --- a/test/docs/transactions.test.js +++ b/test/docs/transactions.test.js @@ -360,7 +360,26 @@ describe('transactions', function() { const test = yield Test.create([{}], { session }).then(res => res[0]); yield test.save(); // throws DocumentNotFoundError })); - yield session.endSession(); + session.endSession(); + }); + }); + + it('correct `isNew` after abort (gh-8852)', function() { + return co(function*() { + const schema = Schema({ name: String }); + + const Test = db.model('gh8852', schema); + + yield Test.createCollection(); + const doc = new Test({ name: 'foo' }); + yield db. + transaction(session => co(function*() { + yield doc.save({ session }); + assert.ok(!doc.isNew); + throw new Error('Oops'); + })). + catch(err => assert.equal(err.message, 'Oops')); + assert.ok(doc.isNew); }); }); });