From ec619000a0db61a8d8ea638938176e0b704e4acb Mon Sep 17 00:00:00 2001 From: Valeri Karpov Date: Thu, 11 Jul 2024 12:15:58 -0400 Subject: [PATCH] feat: add transactionAsyncLocalStorage option to opt in to automatically setting session on all transactions Backport #14583 to 7.x Re: #13889 --- docs/transactions.md | 31 ++++++++++++++++++--- lib/aggregate.js | 5 ++++ lib/connection.js | 13 ++++++--- lib/index.js | 16 +++++++++-- lib/model.js | 7 +++++ lib/query.js | 5 ++++ lib/validoptions.js | 1 + test/docs/transactions.test.js | 49 ++++++++++++++++++++++++++++++++++ types/mongooseoptions.d.ts | 7 +++++ 9 files changed, 126 insertions(+), 8 deletions(-) diff --git a/docs/transactions.md b/docs/transactions.md index 901282dac44..4251cd5d017 100644 --- a/docs/transactions.md +++ b/docs/transactions.md @@ -1,8 +1,6 @@ # Transactions in Mongoose -[Transactions](https://www.mongodb.com/transactions) are new in MongoDB -4.0 and Mongoose 5.2.0. Transactions let you execute multiple operations -in isolation and potentially undo all the operations if one of them fails. +[Transactions](https://www.mongodb.com/transactions) let you execute multiple operations in isolation and potentially undo all the operations if one of them fails. This guide will get you started using transactions with Mongoose.

Getting Started with Transactions

@@ -86,6 +84,33 @@ Below is an example of executing an aggregation within a transaction. [require:transactions.*aggregate] ``` +

Using AsyncLocalStorage

+ +One major pain point with transactions in Mongoose is that you need to remember to set the `session` option on every operation. +If you don't, your operation will execute outside of the transaction. +Mongoose 8.4 is able to set the `session` operation on all operations within a `Connection.prototype.transaction()` executor function using Node's [AsyncLocalStorage API](https://nodejs.org/api/async_context.html#class-asynclocalstorage). +Set the `transactionAsyncLocalStorage` option using `mongoose.set('transactionAsyncLocalStorage', true)` to enable this feature. + +```javascript +mongoose.set('transactionAsyncLocalStorage', true); + +const Test = mongoose.model('Test', mongoose.Schema({ name: String })); + +const doc = new Test({ name: 'test' }); + +// Save a new doc in a transaction that aborts +await connection.transaction(async() => { + await doc.save(); // Notice no session here + throw new Error('Oops'); +}).catch(() => {}); + +// false, `save()` was rolled back +await Test.exists({ _id: doc._id }); +``` + +With `transactionAsyncLocalStorage`, you no longer need to pass sessions to every operation. +Mongoose will add the session by default under the hood. +

Advanced Usage

Advanced users who want more fine-grained control over when they commit or abort transactions diff --git a/lib/aggregate.js b/lib/aggregate.js index 450f9e34d12..0d3ec1bf927 100644 --- a/lib/aggregate.js +++ b/lib/aggregate.js @@ -1022,6 +1022,11 @@ Aggregate.prototype.exec = async function exec() { applyGlobalMaxTimeMS(this.options, model); applyGlobalDiskUse(this.options, model); + const asyncLocalStorage = this.model()?.db?.base.transactionAsyncLocalStorage?.getStore(); + if (!this.options.hasOwnProperty('session') && asyncLocalStorage?.session != null) { + this.options.session = asyncLocalStorage.session; + } + if (this.options && this.options.cursor) { return new AggregationCursor(this); } diff --git a/lib/connection.js b/lib/connection.js index c116a3cde32..bec66623d0b 100644 --- a/lib/connection.js +++ b/lib/connection.js @@ -517,7 +517,7 @@ Connection.prototype.startSession = async function startSession(options) { Connection.prototype.transaction = function transaction(fn, options) { return this.startSession().then(session => { session[sessionNewDocuments] = new Map(); - return session.withTransaction(() => _wrapUserTransaction(fn, session), options). + return session.withTransaction(() => _wrapUserTransaction(fn, session, this.base), options). then(res => { delete session[sessionNewDocuments]; return res; @@ -536,9 +536,16 @@ Connection.prototype.transaction = function transaction(fn, options) { * Reset document state in between transaction retries re: gh-13698 */ -async function _wrapUserTransaction(fn, session) { +async function _wrapUserTransaction(fn, session, mongoose) { try { - const res = await fn(session); + const res = mongoose.transactionAsyncLocalStorage == null + ? await fn(session) + : await new Promise(resolve => { + mongoose.transactionAsyncLocalStorage.run( + { session }, + () => resolve(fn(session)) + ); + }); return res; } catch (err) { _resetSessionDocuments(session); diff --git a/lib/index.js b/lib/index.js index 38197ec50ae..5b716293a4c 100644 --- a/lib/index.js +++ b/lib/index.js @@ -40,6 +40,8 @@ require('./helpers/printJestWarning'); const objectIdHexRegexp = /^[0-9A-Fa-f]{24}$/; +const { AsyncLocalStorage } = require('node:async_hooks'); + /** * Mongoose constructor. * @@ -102,6 +104,10 @@ function Mongoose(options) { } this.Schema.prototype.base = this; + if (options?.transactionAsyncLocalStorage) { + this.transactionAsyncLocalStorage = new AsyncLocalStorage(); + } + Object.defineProperty(this, 'plugins', { configurable: false, enumerable: true, @@ -258,7 +264,7 @@ Mongoose.prototype.set = function(key, value) { if (optionKey === 'objectIdGetter') { if (optionValue) { - Object.defineProperty(mongoose.Types.ObjectId.prototype, '_id', { + Object.defineProperty(_mongoose.Types.ObjectId.prototype, '_id', { enumerable: false, configurable: true, get: function() { @@ -266,7 +272,13 @@ Mongoose.prototype.set = function(key, value) { } }); } else { - delete mongoose.Types.ObjectId.prototype._id; + delete _mongoose.Types.ObjectId.prototype._id; + } + } else if (optionKey === 'transactionAsyncLocalStorage') { + if (optionValue && !_mongoose.transactionAsyncLocalStorage) { + _mongoose.transactionAsyncLocalStorage = new AsyncLocalStorage(); + } else if (!optionValue && _mongoose.transactionAsyncLocalStorage) { + delete _mongoose.transactionAsyncLocalStorage; } } } diff --git a/lib/model.js b/lib/model.js index 3aba07e3563..96b0c78398a 100644 --- a/lib/model.js +++ b/lib/model.js @@ -288,8 +288,11 @@ Model.prototype.$__handleSave = function(options, callback) { } const session = this.$session(); + const asyncLocalStorage = this[modelDbSymbol].base.transactionAsyncLocalStorage?.getStore(); if (!saveOptions.hasOwnProperty('session') && session != null) { saveOptions.session = session; + } else if (asyncLocalStorage?.session != null) { + saveOptions.session = asyncLocalStorage.session; } if (this.$isNew) { @@ -3463,6 +3466,10 @@ Model.bulkWrite = async function bulkWrite(ops, options) { const ordered = options.ordered == null ? true : options.ordered; const validations = ops.map(op => castBulkWrite(this, op, options)); + const asyncLocalStorage = this.db.base.transactionAsyncLocalStorage?.getStore(); + if ((!options || !options.hasOwnProperty('session')) && asyncLocalStorage?.session != null) { + options = { ...options, session: asyncLocalStorage.session }; + } return new Promise((resolve, reject) => { if (ordered) { diff --git a/lib/query.js b/lib/query.js index c4a413fad1f..3f6fc76fb89 100644 --- a/lib/query.js +++ b/lib/query.js @@ -1980,6 +1980,11 @@ Query.prototype._optionsForExec = function(model) { // Apply schema-level `writeConcern` option applyWriteConcern(model.schema, options); + const asyncLocalStorage = this.model?.db?.base.transactionAsyncLocalStorage?.getStore(); + if (!this.options.hasOwnProperty('session') && asyncLocalStorage?.session != null) { + options.session = asyncLocalStorage.session; + } + const readPreference = model && model.schema && model.schema.options && diff --git a/lib/validoptions.js b/lib/validoptions.js index af4e116deec..ee087337bbf 100644 --- a/lib/validoptions.js +++ b/lib/validoptions.js @@ -31,6 +31,7 @@ const VALID_OPTIONS = Object.freeze([ 'strictQuery', 'toJSON', 'toObject', + 'transactionAsyncLocalStorage', 'translateAliases' ]); diff --git a/test/docs/transactions.test.js b/test/docs/transactions.test.js index d196fab7180..bde95f4b176 100644 --- a/test/docs/transactions.test.js +++ b/test/docs/transactions.test.js @@ -421,4 +421,53 @@ describe('transactions', function() { assert.equal(i, 3); }); + + describe('transactionAsyncLocalStorage option', function() { + let m; + before(async function() { + m = new mongoose.Mongoose(); + m.set('transactionAsyncLocalStorage', true); + + await m.connect(start.uri); + }); + + after(async function() { + await m.disconnect(); + }); + + it('transaction() sets `session` by default if transactionAsyncLocalStorage option is set', async function() { + const Test = m.model('Test', m.Schema({ name: String })); + + await Test.createCollection(); + await Test.deleteMany({}); + + const doc = new Test({ name: 'test_transactionAsyncLocalStorage' }); + await assert.rejects( + () => m.connection.transaction(async() => { + await doc.save(); + + await Test.updateOne({ name: 'foo' }, { name: 'foo' }, { upsert: true }); + + let docs = await Test.aggregate([{ $match: { _id: doc._id } }]); + assert.equal(docs.length, 1); + + docs = await Test.find({ _id: doc._id }); + assert.equal(docs.length, 1); + + docs = await async function test() { + return await Test.findOne({ _id: doc._id }); + }(); + assert.equal(doc.name, 'test_transactionAsyncLocalStorage'); + + throw new Error('Oops!'); + }), + /Oops!/ + ); + let exists = await Test.exists({ _id: doc._id }); + assert.ok(!exists); + + exists = await Test.exists({ name: 'foo' }); + assert.ok(!exists); + }); + }); }); diff --git a/types/mongooseoptions.d.ts b/types/mongooseoptions.d.ts index 7fec10b208f..9c35ab8222b 100644 --- a/types/mongooseoptions.d.ts +++ b/types/mongooseoptions.d.ts @@ -203,6 +203,13 @@ declare module 'mongoose' { */ toObject?: ToObjectOptions; + /** + * Set to true to make Mongoose use Node.js' built-in AsyncLocalStorage (Node >= 16.0.0) + * to set `session` option on all operations within a `connection.transaction(fn)` call + * by default. Defaults to false. + */ + transactionAsyncLocalStorage?: boolean; + /** * If `true`, convert any aliases in filter, projection, update, and distinct * to their database property names. Defaults to false.