feat: add transactionAsyncLocalStorage option to opt in to automatically setting session on all transactions

Backport #14583 to 7.x
Re: #13889

Author: vkarpov15

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.
 
 <h2 id="getting-started-with-transactions"><a href="#getting-started-with-transactions">Getting Started with Transactions</a></h2>
@@ -86,6 +84,33 @@ Below is an example of executing an aggregation within a transaction.
 [require:transactions.*aggregate]
 ```
 
+<h2 id="asynclocalstorage"><a href="#asynclocalstorage">Using AsyncLocalStorage</a></h2>
+
+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.
+
 <h2 id="advanced-usage"><a href="#advanced-usage">Advanced Usage</a></h2>
 
 Advanced users who want more fine-grained control over when they commit or abort transactions

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);
   }

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);

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,15 +264,21 @@ 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() {
             return this;
           }
         });
       } 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;
       }
     }
   }

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) {

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 &&

lib/validoptions.js

@@ -31,6 +31,7 @@ const VALID_OPTIONS = Object.freeze([
   'strictQuery',
   'toJSON',
   'toObject',
+  'transactionAsyncLocalStorage',
   'translateAliases'
 ]);
 

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);
+    });
+  });
 });

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.