feat(shell-api): add Session.withTransaction MONGOSH-1001 (#1127)

Author: addaleax

packages/i18n/src/locales/en_US.ts

@@ -1673,6 +1673,11 @@ const translations: Catalog = {
               link: 'https://docs.mongodb.com/manual/reference/method/Session.abortTransaction/#Session.abortTransaction',
               description: 'Aborts the session’s transaction.'
             },
+            withTransaction: {
+              link: 'https://docs.mongodb.com/manual/reference/method/Session.abortTransaction/#Session.withTransaction',
+              description: 'Run a function within a transaction context.',
+              example: 'session.withTransaction(() => { session.getDatabase("test").test.insertOne({ doc: 1 }); })'
+            },
             advanceClusterTime: {
               description: 'Advances the clusterTime for a Session to the provided clusterTime.'
             }

packages/shell-api/src/session.spec.ts

@@ -138,6 +138,11 @@ describe('Session', () => {
       await session.abortTransaction();
       expect(serviceProviderSession.abortTransaction).to.have.been.calledOnceWith();
     });
+    it('withTransaction', async() => {
+      serviceProviderSession.withTransaction.resolves();
+      await session.withTransaction(() => {});
+      expect(serviceProviderSession.withTransaction).to.have.been.calledOnce;
+    });
   });
   describe('integration', () => {
     const [ srv0 ] = startTestCluster(['--replicaset']);
@@ -286,6 +291,42 @@ describe('Session', () => {
         await session.abortTransaction();
         expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
       });
+      it('can run withTransaction in the success case', async() => {
+        const doc = { value: 'test', count: 0 };
+        const testColl = mongo.getDB(databaseName).getCollection('coll');
+        await testColl.drop();
+        await testColl.insertOne(doc);
+        expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
+        session = mongo.startSession();
+        await session.withTransaction(async() => {
+          const sessionColl = session.getDatabase(databaseName).getCollection('coll');
+          expect((await sessionColl.updateOne(
+            { value: 'test' },
+            { $inc: { count: 1 } }
+          )).acknowledged).to.be.true;
+          expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
+        });
+        expect((await testColl.findOne({ value: 'test' })).count).to.equal(1);
+      });
+      it('can run withTransaction in the failure case', async() => {
+        const doc = { value: 'test', count: 0 };
+        const testColl = mongo.getDB(databaseName).getCollection('coll');
+        await testColl.drop();
+        await testColl.insertOne(doc);
+        expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
+        session = mongo.startSession();
+        const { err } = await session.withTransaction(async() => {
+          const sessionColl = session.getDatabase(databaseName).getCollection('coll');
+          expect((await sessionColl.updateOne(
+            { value: 'test' },
+            { $inc: { count: 1 } }
+          )).acknowledged).to.be.true;
+          expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
+          throw new Error('fails');
+        }).catch(err => ({ err }));
+        expect(err.message).to.equal('fails');
+        expect((await testColl.findOne({ value: 'test' })).count).to.equal(0);
+      });
     });
     describe('after resetting connection will error with expired session', () => {
       skipIfApiStrict();

packages/shell-api/src/session.ts

@@ -21,7 +21,7 @@ import { CommonErrors, MongoshInvalidInputError } from '@mongosh/errors';
 import { assertArgsDefinedType, isValidDatabaseName } from './helpers';
 
 @shellApiClassDefault
-@classPlatforms([ ReplPlatform.CLI ] )
+@classPlatforms([ ReplPlatform.CLI ])
 export default class Session extends ShellApiWithMongoClass {
   public id: ServerSessionId | undefined;
   public _session: ClientSession;
@@ -100,4 +100,12 @@ export default class Session extends ShellApiWithMongoClass {
   async abortTransaction(): Promise<Document> {
     return await this._session.abortTransaction();
   }
+
+  @returnsPromise
+  async withTransaction<T extends(...args: any) => any>(fn: T, options: TransactionOptions = {}): Promise<ReturnType<T>> {
+    assertArgsDefinedType([fn, options], ['function', [undefined, 'object']]);
+    // The driver doesn't automatically ensure that fn is an async
+    // function/convert its return type to a Promise, so we do that here.
+    return await this._session.withTransaction(async() => await fn(), options);
+  }
 }