diff --git a/docs/developer-guide/transactions.md b/docs/developer-guide/transactions.md new file mode 100644 index 000000000..7657d247a --- /dev/null +++ b/docs/developer-guide/transactions.md @@ -0,0 +1,243 @@ +--- +title: MongoDB Transactional Support +audience: Technical +created_by: minottic +created_on: 2026-08-04 +--- + +# MongoDB Transactional Support + +> ⚠️ **Warning:** MongoDB transactions require the server to be running as a +> **replica set** (or a sharded cluster) — they are not supported on a +> standalone `mongod`, which is what most local/dev setups use by default. +> See MongoDB's docs on [transactions](https://www.mongodb.com/docs/manual/core/transactions/) +> and on [deploying a replica set](https://www.mongodb.com/docs/manual/tutorial/deploy-replica-set/) +> for details. `MongoTransactionService` detects this automatically and +> falls back to running without a transaction rather than failing — see +> [Deployments without replica sets](#deployments-without-replica-sets) below. + +## Overview + +This is the framework for running MongoDB operations atomically inside a +transaction. It's introduced here without being applied to any schema or +service yet — later PRs will opt specific models and methods into it. This +document explains what's available and what a PR needs to do to actually use +it. + +The pieces: + +- `MongoTransactionService` (`src/common/services/mongo-transaction.service.ts`) — + the core primitive. Runs a function inside a MongoDB session/transaction. +- `@Transactional()` (`src/common/decorators/transactional.decorator.ts`) — + a method decorator that wraps a whole method body in + `MongoTransactionService.run()`, so you don't have to call it by hand. +- `sessionPlugin` (`src/common/mongoose/plugins/session.plugin.ts`) — a + mongoose schema plugin that automatically attaches the active transaction + session to every query, aggregate, and save on that schema, so callers + don't need to pass `{ session }` explicitly. **Must be registered on a + schema for that schema's models to participate in transactions + automatically** — see [Registering the plugin on a schema](#registering-the-plugin-on-a-schema). +- `session-context.util.ts` — the underlying `AsyncLocalStorage` plumbing + that makes the session "ambient" (see below). You shouldn't need to use + this directly. + +## Core concept: the ambient session + +Once inside a transaction, the active `ClientSession` is made "ambient" — +readable via `getCurrentSession()` from anywhere in the same call chain, +without being passed as an argument. `sessionPlugin` uses this to attach the +session to a query automatically. This is what lets you write: + +```ts +return this.mongoTransactionService.run(async () => { + const [dataset] = await this.datasetModel.create([dto]); + await this.datablocksService.createBlocks(dataset); + return dataset; +}); +``` + +instead of manually threading `{ session }` through every call, as long as +`DatasetSchema` and the datablock schema have `sessionPlugin` registered. + +An explicit `{ session }` passed to any individual query always wins over +the ambient one — including `session: null`, which is respected as "opt this +one query out of the transaction" rather than being treated as "no session +set." See `attachAmbientSession` in `session.plugin.ts` for the exact rule. + +## How to use it + +### 1. Registering the plugin on a schema + +For `sessionPlugin` to have any effect, the model's schema must register it: + +```ts +import { sessionPlugin } from "src/common/mongoose/plugins/session.plugin"; + +// ...after SchemaFactory.createForClass(...) +DatasetSchema.plugin(sessionPlugin); +``` + +Registering the plugin is a per-schema, opt-in choice — it lets you decide +selectively which models pick up the ambient session automatically. If you +want more control over a given model, and would rather avoid automatic +session pickup for it, simply don't register the plugin on that schema: +queries on it will then only join a transaction when you pass `{ session }` +explicitly. It currently isn't registered on any schema in this codebase. + +### 2. Injecting `MongoTransactionService` + +Both the decorator and calling `run()` directly require the owning class to +inject `MongoTransactionService` via its constructor: + +```ts +@Injectable() +export class DatasetsService { + constructor( + private readonly mongoTransactionService: MongoTransactionService, + @InjectModel(Dataset.name) private datasetModel: Model, + ) {} +} +``` + +### 3. Using `@Transactional()` + +For the common case — wrap this whole method in a transaction — annotate +the method. The property **must** be named `mongoTransactionService`; the +decorator reads it off `this` at call time: + +```ts +@Transactional() +async updateOne(id: string, update: UpdateDatasetDto) { + const dataset = await this.datasetModel.findOneAndUpdate({ _id: id }, update); + await this.recomputeCounts(dataset); // also joins the same transaction + return dataset; +} +``` + +The decorated method must be `async` and return a `Promise`. If the owning +class doesn't inject `MongoTransactionService` as `this.mongoTransactionService`, +calling the method throws a clear error naming the class and the fix needed, +rather than a generic "cannot read properties of undefined." + +### 4. Calling `MongoTransactionService.run()` directly + +Use this when you need more control than a bare decorator gives you, or when +you're not inside a class that has `@Transactional()` available. + +`fn` always receives the session as its argument, in addition to it being +made ambient — use whichever is more convenient at each call site. + +**Passing the session explicitly** is necessary in two situations, even on +schemas where `sessionPlugin` is already attaching it to your queries +automatically: + +- Anything `sessionPlugin` can't reach — it only hooks mongoose + query/aggregate/document middleware, not raw driver-level collection + methods like `bulkWrite`: + + ```ts + return this.mongoTransactionService.run(async (session) => { + return this.datasetModel.collection.bulkWrite(operations, { session }); + }); + ``` + +- Calling a method directly on the session itself, rather than attaching it + to a query — `sessionPlugin` only ever attaches the session *to* queries, + it never exposes the session object for you to act on. For example, + ending the transaction early based on a business check instead of + throwing an error — the driver explicitly supports this: if `fn` calls + `session.abortTransaction()` itself, `withTransaction` detects that and + returns without attempting to commit: + + ```ts + return this.mongoTransactionService.run(async (session) => { + const [dataset] = await this.datasetModel.create([dto]); + if (!isStillValid(dataset)) { + await session.abortTransaction(); + return null; + } + return dataset; + }); + ``` + +**Without passing the session** works for everything `sessionPlugin` already +covers — plain mongoose calls just pick up the ambient session on their own, +the same way they would inside a `@Transactional()` method: + +```ts +return this.mongoTransactionService.run(async () => { + const [dataset] = await this.datasetModel.create([dto]); + await this.datablocksService.createBlocks(dataset); + return dataset; +}); +``` + +### 5. Explicit mode (`{ ambient: false }`) + +By default, `run()` makes the session ambient (as described above). Pass +`{ ambient: false }` to disable that — nothing will be attached +automatically, and every query that should join the transaction must be +given `{ session }` explicitly: + +```ts +return this.mongoTransactionService.run( + async (session) => { + const [dataset] = await this.datasetModel.create([dto], { session }); + await this.createBlocks(dataset, session); // threaded to a sibling + return dataset; + }, + { ambient: false }, +); +``` + +This is useful when you want the transaction boundary to be visible at every +call site rather than implicit, or when working with a schema that doesn't +have `sessionPlugin` registered. + +### Nesting + +Calling `run()` (or a `@Transactional()` method) from inside another active +`run()` call joins the existing transaction instead of starting a second, +independent one — only the outermost call actually opens a session and +commits/aborts it. This works automatically in the default (ambient) mode. +In explicit mode, nesting isn't auto-detected — pass the session you +received down as `existingSession` to join it: + +```ts +return this.mongoTransactionService.run( + async (session) => + this.mongoTransactionService.run(fn, { + ambient: false, + existingSession: session, + }), + { ambient: false }, +); +``` + +Don't mix ambient and explicit calls inside one another without threading +the session through by hand — the two modes can't detect each other, so +you'd end up with two separate, possibly conflicting transactions instead of +one. + +## Deployments without replica sets + +`MongoTransactionService.run()` detects when the deployment doesn't support +transactions (e.g. a standalone `mongod`, common in local/CI) and falls back +to running `fn` without a session, rather than throwing. The result is +cached on the service instance after the first failed attempt, so later +calls skip straight to the fallback instead of repeatedly starting and +failing a transaction. + +## Watch out for un-awaited promises + +Because the session is scoped to the lifetime of the wrapped function's +returned promise, an async call started inside a `@Transactional()` method +(or a `run()` callback) but not `await`-ed can still be in flight when the +transaction commits and the session closes — leading to a driver error, or a +write racing the commit and silently getting lost. Every async call inside a +transactional method must be awaited. This project enforces +`@typescript-eslint/no-floating-promises` and +`@typescript-eslint/no-misused-promises` project-wide, which catches the +common accidental cases (a forgotten `await`, `array.forEach(async ...)`) — +but not a deliberately `void`-marked call, or a callback registered through +a loosely-typed API. Those still require care in review. diff --git a/src/app.module.ts b/src/app.module.ts index 1295cbfae..314a9a299 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -1,6 +1,7 @@ import { Module, NestModule } from "@nestjs/common"; import { MongooseModule } from "@nestjs/mongoose"; import { RequestContextModule } from "./common/modules/request-context.module"; +import { MongoTransactionModule } from "./common/modules/mongo-transaction.module"; import { DatasetsModule } from "./datasets/datasets.module"; import { AuthModule } from "./auth/auth.module"; import { UsersModule } from "./users/users.module"; @@ -162,6 +163,7 @@ import { ThrottlerModule } from "@nestjs/throttler"; AdminModule, HealthModule, RequestContextModule, + MongoTransactionModule, HistoryModule, ConditionalModule.registerWhen( MaskSensitiveDataInterceptorModule, diff --git a/src/common/decorators/transactional.decorator.spec.ts b/src/common/decorators/transactional.decorator.spec.ts new file mode 100644 index 000000000..c40ae9453 --- /dev/null +++ b/src/common/decorators/transactional.decorator.spec.ts @@ -0,0 +1,88 @@ +import { Transactional } from "./transactional.decorator"; +import { MongoTransactionService } from "../services/mongo-transaction.service"; +import { getCurrentSession } from "../utils/session-context.util"; + +describe("Transactional", () => { + class TestService { + mongoTransactionService = { + run: jest.fn().mockImplementation((fn) => fn()), + }; + + @Transactional() + async withArgs(a: string, b: string) { + return { a, b }; + } + } + + let service: TestService; + + beforeEach(() => { + service = new TestService(); + }); + + it("delegates to mongoTransactionService.run and forwards all arguments unchanged", async () => { + const result = await service.withArgs("x", "y"); + + expect(service.mongoTransactionService.run).toHaveBeenCalledTimes(1); + expect(result).toEqual({ a: "x", b: "y" }); + }); + + it("throws a clear error when the owning class doesn't inject MongoTransactionService", async () => { + class MissingService { + @Transactional() + async doThing() { + return "done"; + } + } + + await expect(new MissingService().doThing()).rejects.toThrow( + "@Transactional() requires MissingService to inject MongoTransactionService as this.mongoTransactionService", + ); + }); + + describe("nesting on top of the real MongoTransactionService", () => { + const mockSession = { + withTransaction: jest.fn().mockImplementation((fn) => fn()), + endSession: jest.fn().mockResolvedValue(undefined), + }; + const connection = { + startSession: jest.fn().mockResolvedValue(mockSession), + }; + + class RealTransactionService { + mongoTransactionService = new MongoTransactionService( + connection as never, + ); + sessionsSeen: unknown[] = []; + + @Transactional() + async outer() { + this.sessionsSeen.push(getCurrentSession()); + return this.inner(); + } + + @Transactional() + async inner() { + this.sessionsSeen.push(getCurrentSession()); + return "inner-done"; + } + } + + beforeEach(() => { + jest.clearAllMocks(); + mockSession.withTransaction.mockImplementation((fn) => fn()); + connection.startSession.mockResolvedValue(mockSession); + }); + + it("joins the outer transaction instead of nesting a second one", async () => { + const realService = new RealTransactionService(); + + const result = await realService.outer(); + + expect(result).toBe("inner-done"); + expect(connection.startSession).toHaveBeenCalledTimes(1); + expect(mockSession.endSession).toHaveBeenCalledTimes(1); + expect(realService.sessionsSeen).toEqual([mockSession, mockSession]); + }); + }); +}); diff --git a/src/common/decorators/transactional.decorator.ts b/src/common/decorators/transactional.decorator.ts new file mode 100644 index 000000000..fcef24cc6 --- /dev/null +++ b/src/common/decorators/transactional.decorator.ts @@ -0,0 +1,43 @@ +import { MongoTransactionService } from "../services/mongo-transaction.service"; + +interface WithMongoTransactionService { + mongoTransactionService: MongoTransactionService; +} + +/** + * Runs the decorated method inside a MongoDB transaction via + * `MongoTransactionService.run`, joining an already-active transaction + * instead of nesting a new one when called from another `@Transactional()` + * method. + * + * The session is ambient, read via `getCurrentSession()` by whatever the + * decorated method calls. The method must be async and return a Promise. + * The owning class must inject `MongoTransactionService` as + * `this.mongoTransactionService`. + */ +export function Transactional(): MethodDecorator { + return function ( + _target: object, + _propertyKey: string | symbol, + descriptor: PropertyDescriptor, + ) { + const originalMethod = descriptor.value; + + descriptor.value = async function ( + this: WithMongoTransactionService, + ...args: unknown[] + ) { + if (!this.mongoTransactionService) { + throw new Error( + `@Transactional() requires ${this.constructor.name} to inject ` + + "MongoTransactionService as this.mongoTransactionService", + ); + } + return this.mongoTransactionService.run(() => + originalMethod.apply(this, args), + ); + }; + + return descriptor; + }; +} diff --git a/src/common/modules/mongo-transaction.module.ts b/src/common/modules/mongo-transaction.module.ts new file mode 100644 index 000000000..f29b278fd --- /dev/null +++ b/src/common/modules/mongo-transaction.module.ts @@ -0,0 +1,9 @@ +import { Global, Module } from "@nestjs/common"; +import { MongoTransactionService } from "../services/mongo-transaction.service"; + +@Global() +@Module({ + providers: [MongoTransactionService], + exports: [MongoTransactionService], +}) +export class MongoTransactionModule {} diff --git a/src/common/mongoose/plugins/session.plugin.spec.ts b/src/common/mongoose/plugins/session.plugin.spec.ts new file mode 100644 index 000000000..c3f1b4ede --- /dev/null +++ b/src/common/mongoose/plugins/session.plugin.spec.ts @@ -0,0 +1,150 @@ +import { Schema } from "mongoose"; +import { attachAmbientSession, sessionPlugin } from "./session.plugin"; +import { runWithSession } from "../../utils/session-context.util"; + +describe("attachAmbientSession", () => { + const fakeSession = { id: "s1" } as never; + + it("does nothing outside of a transaction", () => { + const query = { + session: jest.fn(), + getOptions: jest.fn().mockReturnValue({}), + }; + + attachAmbientSession(query); + + expect(query.session).not.toHaveBeenCalled(); + }); + + it("attaches the ambient session to a query with no session of its own", async () => { + const query = { + session: jest.fn(), + getOptions: jest.fn().mockReturnValue({}), + }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(query); + }); + + expect(query.session).toHaveBeenCalledWith(fakeSession); + }); + + it("does not override a query that already has its own session", async () => { + const ownSession = { id: "own" } as never; + const query = { + session: jest.fn(), + getOptions: jest.fn().mockReturnValue({ session: ownSession }), + }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(query); + }); + + expect(query.session).not.toHaveBeenCalled(); + }); + + it("does not override a query explicitly opted out of any session with session: null", async () => { + const query = { + session: jest.fn(), + getOptions: jest.fn().mockReturnValue({ session: null }), + }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(query); + }); + + expect(query.session).not.toHaveBeenCalled(); + }); + + it("attaches the ambient session to an aggregate with no session of its own", async () => { + const aggregate = { session: jest.fn(), options: {} }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(aggregate); + }); + + expect(aggregate.session).toHaveBeenCalledWith(fakeSession); + }); + + it("does not override an aggregate that already has its own session", async () => { + const ownSession = { id: "own" } as never; + const aggregate = { session: jest.fn(), options: { session: ownSession } }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(aggregate); + }); + + expect(aggregate.session).not.toHaveBeenCalled(); + }); + + it("does not override an aggregate explicitly opted out of any session with session: null", async () => { + const aggregate = { session: jest.fn(), options: { session: null } }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(aggregate); + }); + + expect(aggregate.session).not.toHaveBeenCalled(); + }); + + it("attaches the ambient session to a document with no session of its own", async () => { + const doc = { $session: jest.fn().mockReturnValue(undefined) }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(doc); + }); + + expect(doc.$session).toHaveBeenCalledWith(fakeSession); + }); + + it("does not override a document that already has its own session", async () => { + const ownSession = { id: "own" }; + const doc = { $session: jest.fn().mockReturnValue(ownSession) }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(doc); + }); + + expect(doc.$session).toHaveBeenCalledTimes(1); + }); + + it("does not override a document explicitly opted out of any session with $session(null)", async () => { + const doc = { $session: jest.fn().mockReturnValue(null) }; + + await runWithSession(fakeSession, async () => { + attachAmbientSession(doc); + }); + + expect(doc.$session).toHaveBeenCalledTimes(1); + }); +}); + +describe("sessionPlugin", () => { + it("registers pre-hooks for queries, aggregate, and save", () => { + const schema = new Schema({ name: String }); + const preSpy = jest.spyOn(schema, "pre"); + + sessionPlugin(schema); + + const registeredOps = preSpy.mock.calls.flatMap((call) => call[0]); + expect(registeredOps).toEqual( + expect.arrayContaining([ + "countDocuments", + "distinct", + "estimatedDocumentCount", + "find", + "findOne", + "findOneAndReplace", + "findOneAndUpdate", + "replaceOne", + "updateMany", + "updateOne", + "deleteMany", + "deleteOne", + "findOneAndDelete", + "aggregate", + "save", + ]), + ); + }); +}); diff --git a/src/common/mongoose/plugins/session.plugin.ts b/src/common/mongoose/plugins/session.plugin.ts new file mode 100644 index 000000000..9fa89f5b9 --- /dev/null +++ b/src/common/mongoose/plugins/session.plugin.ts @@ -0,0 +1,109 @@ +import { ClientSession, Schema } from "mongoose"; +import { getCurrentSession } from "../../utils/session-context.util"; + +interface SessionCapableQuery { + session(session: ClientSession | null): unknown; + getOptions(): { session?: ClientSession | null }; +} + +interface SessionCapableAggregate { + session(session: ClientSession | null): unknown; + options?: { session?: ClientSession | null }; +} + +interface SessionCapableDocument { + $session(session?: ClientSession | null): ClientSession | null | undefined; +} + +/** + * Attaches the ambient session to a mongoose query, aggregate, or + * document context, unless it already has a session of its own — + * explicit `session: null` counts as "of its own" too, so this checks + * for the key's presence rather than truthiness. + * @param context The query, aggregate, or document to attach to. + */ +export function attachAmbientSession( + context: + SessionCapableQuery | SessionCapableAggregate | SessionCapableDocument, +): void { + const session = getCurrentSession(); + if (!session) return; + + if ("$session" in context) { + if (context.$session() === undefined) context.$session(session); + return; + } + + if ("getOptions" in context) { + if (!("session" in context.getOptions())) context.session(session); + return; + } + + if (!context.options || !("session" in context.options)) { + context.session(session); + } +} + +// Mirrors mongoose's own `queryOperations` list (lib/constants.js) — every +// operation type mongoose recognizes as query middleware. +type QueryMiddlewareOp = + | "countDocuments" + | "distinct" + | "estimatedDocumentCount" + | "find" + | "findOne" + | "findOneAndReplace" + | "findOneAndUpdate" + | "replaceOne" + | "updateMany" + | "updateOne" + | "deleteMany" + | "deleteOne" + | "findOneAndDelete"; + +const QUERY_MIDDLEWARE_OPS: QueryMiddlewareOp[] = [ + "countDocuments", + "distinct", + "estimatedDocumentCount", + "find", + "findOne", + "findOneAndReplace", + "findOneAndUpdate", + "replaceOne", + "updateMany", + "updateOne", + "deleteMany", + "deleteOne", + "findOneAndDelete", +]; + +/** + * Mongoose plugin that attaches the active transaction session to every + * query, aggregate, and save on the schema, unless that operation already + * specifies its own. + * @param schema The schema to attach the hooks to. + * @example + * DatablockSchema.plugin(sessionPlugin); + */ +export function sessionPlugin(schema: Schema) { + schema.pre( + QUERY_MIDDLEWARE_OPS, + function (this: SessionCapableQuery, next: () => void) { + attachAmbientSession(this); + next(); + }, + ); + + schema.pre( + "aggregate", + function (this: SessionCapableAggregate, next: () => void) { + attachAmbientSession(this); + next(); + }, + ); + + schema.pre("save", function (this: SessionCapableDocument, next: () => void) { + attachAmbientSession(this); + next(); + }); +} diff --git a/src/common/services/mongo-transaction.service.spec.ts b/src/common/services/mongo-transaction.service.spec.ts new file mode 100644 index 000000000..948e8a52f --- /dev/null +++ b/src/common/services/mongo-transaction.service.spec.ts @@ -0,0 +1,195 @@ +import { MongoServerError } from "mongodb"; +import { MongoTransactionService } from "./mongo-transaction.service"; +import { getCurrentSession } from "../utils/session-context.util"; + +describe("MongoTransactionService", () => { + const mockSession = { + withTransaction: jest.fn().mockImplementation((fn) => fn()), + endSession: jest.fn().mockResolvedValue(undefined), + }; + + const connection = { + startSession: jest.fn().mockResolvedValue(mockSession), + }; + + let service: MongoTransactionService; + + beforeEach(() => { + jest.clearAllMocks(); + mockSession.withTransaction.mockImplementation((fn) => fn()); + connection.startSession.mockResolvedValue(mockSession); + service = new MongoTransactionService(connection as never); + }); + + it("runs fn with the session available via getCurrentSession", async () => { + const result = await service.run(async () => getCurrentSession()); + + expect(result).toBe(mockSession); + expect(mockSession.endSession).toHaveBeenCalled(); + }); + + it("also passes the session explicitly as fn's argument", async () => { + const fn = jest.fn().mockResolvedValue("done"); + + await service.run(fn); + + expect(fn).toHaveBeenCalledWith(mockSession); + }); + + it("does not leak the session outside of run", async () => { + await service.run(async () => getCurrentSession()); + + expect(getCurrentSession()).toBeUndefined(); + }); + + it("falls back to a non-transactional run when transactions are not supported", async () => { + const notSupportedError = new MongoServerError({ + message: + "Transaction numbers are only allowed on a replica set member or mongos", + }); + notSupportedError.code = 20; + mockSession.withTransaction.mockImplementationOnce(() => { + throw notSupportedError; + }); + const fn = jest.fn().mockImplementation(async () => getCurrentSession()); + + const result = await service.run(fn); + + expect(fn).toHaveBeenCalledTimes(1); + expect(fn).toHaveBeenCalledWith(undefined); + expect(result).toBeUndefined(); + expect(mockSession.endSession).toHaveBeenCalled(); + }); + + it("caches unsupported-transactions detection so later calls skip starting a session", async () => { + const notSupportedError = new MongoServerError({ + message: + "Transaction numbers are only allowed on a replica set member or mongos", + }); + notSupportedError.code = 20; + mockSession.withTransaction.mockImplementationOnce(() => { + throw notSupportedError; + }); + + await service.run(async () => "first"); + await service.run(async () => "second"); + + expect(connection.startSession).toHaveBeenCalledTimes(1); + }); + + it("rethrows errors that are not the transactions-not-supported case", async () => { + const otherError = new MongoServerError({ message: "boom" }); + otherError.code = 11000; + mockSession.withTransaction.mockImplementationOnce(() => { + throw otherError; + }); + + await expect(service.run(async () => "unreachable")).rejects.toThrow( + otherError, + ); + expect(mockSession.endSession).toHaveBeenCalled(); + }); + + it("ends the session even when fn rejects", async () => { + await expect( + service.run(async () => { + throw new Error("fn failed"); + }), + ).rejects.toThrow("fn failed"); + expect(mockSession.endSession).toHaveBeenCalled(); + }); + + it("joins an already-active transaction instead of starting a nested one", async () => { + const sessionsSeen: unknown[] = []; + + const result = await service.run(async () => { + sessionsSeen.push(getCurrentSession()); + return service.run(async () => { + sessionsSeen.push(getCurrentSession()); + return "nested-done"; + }); + }); + + expect(result).toBe("nested-done"); + expect(connection.startSession).toHaveBeenCalledTimes(1); + expect(mockSession.endSession).toHaveBeenCalledTimes(1); + expect(sessionsSeen).toEqual([mockSession, mockSession]); + }); + + describe("ambient: false", () => { + it("passes the session as fn's argument without making it ambient", async () => { + const fn = jest.fn().mockImplementation(async () => getCurrentSession()); + + const result = await service.run(fn, { + ambient: false, + }); + + expect(fn).toHaveBeenCalledWith(mockSession); + expect(result).toBeUndefined(); + }); + + it("starts an independent transaction when nested without existingSession", async () => { + const result = await service.run( + async () => + service.run(async () => "nested-done", { + ambient: false, + }), + { ambient: false }, + ); + + expect(result).toBe("nested-done"); + expect(connection.startSession).toHaveBeenCalledTimes(2); + expect(mockSession.endSession).toHaveBeenCalledTimes(2); + }); + + it("joins the given existingSession instead of starting a new transaction", async () => { + const fn = jest.fn().mockResolvedValue("done"); + + const result = await service.run(fn, { + ambient: false, + existingSession: mockSession as never, + }); + + expect(result).toBe("done"); + expect(fn).toHaveBeenCalledWith(mockSession); + expect(connection.startSession).not.toHaveBeenCalled(); + }); + + it("falls back to a non-transactional run when transactions are not supported", async () => { + const notSupportedError = new MongoServerError({ + message: + "Transaction numbers are only allowed on a replica set member or mongos", + }); + notSupportedError.code = 20; + mockSession.withTransaction.mockImplementationOnce(() => { + throw notSupportedError; + }); + const fn = jest.fn().mockResolvedValue("done"); + + const result = await service.run(fn, { + ambient: false, + }); + + expect(fn).toHaveBeenCalledWith(undefined); + expect(result).toBe("done"); + }); + + it("shares the unsupported-transactions cache with the ambient mode", async () => { + const notSupportedError = new MongoServerError({ + message: + "Transaction numbers are only allowed on a replica set member or mongos", + }); + notSupportedError.code = 20; + mockSession.withTransaction.mockImplementationOnce(() => { + throw notSupportedError; + }); + + await service.run(async () => "first"); + const fn = jest.fn().mockResolvedValue("second"); + await service.run(fn, { ambient: false }); + + expect(connection.startSession).toHaveBeenCalledTimes(1); + expect(fn).toHaveBeenCalledWith(undefined); + }); + }); +}); diff --git a/src/common/services/mongo-transaction.service.ts b/src/common/services/mongo-transaction.service.ts new file mode 100644 index 000000000..71c5dcdd5 --- /dev/null +++ b/src/common/services/mongo-transaction.service.ts @@ -0,0 +1,88 @@ +import { Injectable } from "@nestjs/common"; +import { InjectConnection } from "@nestjs/mongoose"; +import { MongoServerError } from "mongodb"; +import { ClientSession, Connection } from "mongoose"; +import { + getCurrentSession, + runWithSession, +} from "../utils/session-context.util"; + +const TRANSACTIONS_NOT_SUPPORTED_CODE = 20; // IllegalOperation: no replica set/mongos + +type RunOptions = + { ambient?: true } | { ambient: false; existingSession?: ClientSession }; + +@Injectable() +export class MongoTransactionService { + private transactionsSupported: boolean | undefined; + + constructor(@InjectConnection() private readonly connection: Connection) {} + + /** + * Runs `fn` inside a MongoDB transaction, falling back to a + * non-transactional run if the deployment doesn't support transactions + * + * By default, the session is both passed to `fn` and made ambient (via + * `getCurrentSession()`/`sessionPlugin`), so nested calls pick it up + * without it being passed explicitly; nesting inside another `run()` + * joins that transaction automatically. With `ambient: false`, only + * `fn`'s argument carries the session — nothing attaches it + * automatically — and nesting must be done by hand via + * `existingSession`. + * @param fn The operation to run. Pass the session if you later + * want to use it inside the closure. + * @param options.ambient Set to `false` to disable automatic + * attachment; see above. + * @param options.existingSession Only used with `ambient: false`: pass + * the session from an enclosing call to join it instead of starting a + * new transaction. + * @returns Whatever `fn` resolves to. + * @example + * return this.mongoTransactionService.run(async () => { + * const [dataset] = await this.datasetModel.create([dto]); + * return dataset; + * }); + * @example + * return this.mongoTransactionService.run( + * async (session) => this.datasetModel.create([dto], { session }), + * { ambient: false }, + * ); + */ + async run( + fn: (session: ClientSession | undefined) => Promise, + options: RunOptions = {}, + ): Promise { + if (options.ambient === false) { + if (options.existingSession) return fn(options.existingSession); + } else { + const currentSession = getCurrentSession(); + if (currentSession) return fn(currentSession); + } + + let session: ClientSession | undefined; + try { + if (this.transactionsSupported === false) return await fn(undefined); + + session = await this.connection.startSession(); + let result!: T; + await session.withTransaction(async () => { + result = + options.ambient === false + ? await fn(session) + : await runWithSession(session as ClientSession, () => fn(session)); + }); + return result; + } catch (error) { + if ( + error instanceof MongoServerError && + error.code === TRANSACTIONS_NOT_SUPPORTED_CODE + ) { + this.transactionsSupported = false; + return await fn(undefined); + } + throw error; + } finally { + if (session) await session.endSession(); + } + } +} diff --git a/src/common/utils/session-context.util.spec.ts b/src/common/utils/session-context.util.spec.ts new file mode 100644 index 000000000..d79bbaeca --- /dev/null +++ b/src/common/utils/session-context.util.spec.ts @@ -0,0 +1,67 @@ +import { getCurrentSession, runWithSession } from "./session-context.util"; + +describe("session-context.util", () => { + it("returns undefined outside any session scope", () => { + expect(getCurrentSession()).toBeUndefined(); + }); + + it("makes the session available inside runWithSession", async () => { + const session = { id: "s1" } as unknown as import("mongoose").ClientSession; + + const seenInside = await runWithSession(session, async () => { + return getCurrentSession(); + }); + + expect(seenInside).toBe(session); + }); + + it("propagates the session across awaits within the scope", async () => { + const session = { id: "s2" } as unknown as import("mongoose").ClientSession; + + const seen = await runWithSession(session, async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + return getCurrentSession(); + }); + + expect(seen).toBe(session); + }); + + it("propagates the session into concurrent branches started inside the scope", async () => { + const session = { id: "s3" } as unknown as import("mongoose").ClientSession; + + const [a, b] = await runWithSession(session, () => + Promise.all([ + (async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + return getCurrentSession(); + })(), + (async () => getCurrentSession())(), + ]), + ); + + expect(a).toBe(session); + expect(b).toBe(session); + }); + + it("does not leak the session after runWithSession resolves", async () => { + const session = { id: "s4" } as unknown as import("mongoose").ClientSession; + + await runWithSession(session, async () => getCurrentSession()); + + expect(getCurrentSession()).toBeUndefined(); + }); + + it("does not leak the session to code that runs concurrently but outside the scope", async () => { + const session = { id: "s5" } as unknown as import("mongoose").ClientSession; + + const outsidePromise = new Promise((resolve) => + setTimeout(() => resolve(getCurrentSession()), 0), + ); + + await runWithSession(session, async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + }); + + expect(await outsidePromise).toBeUndefined(); + }); +}); diff --git a/src/common/utils/session-context.util.ts b/src/common/utils/session-context.util.ts new file mode 100644 index 000000000..90c15e167 --- /dev/null +++ b/src/common/utils/session-context.util.ts @@ -0,0 +1,30 @@ +import { AsyncLocalStorage } from "async_hooks"; +import { ClientSession } from "mongoose"; + +/** + * Holds the currently active MongoDB transaction session, if any. + */ +const asyncLocalStorage = new AsyncLocalStorage(); + +/** + * Runs `fn` with `session` as the current session for its duration (and + * everything it awaits), so `getCurrentSession()` picks it up anywhere in + * that call chain without it being passed explicitly. + * @param session The session to make current for `fn`. + * @param fn The operation to run within the session's scope. + * @returns Whatever `fn` resolves to. + */ +export function runWithSession( + session: ClientSession, + fn: () => Promise, +): Promise { + return asyncLocalStorage.run(session, fn); +} + +/** + * Gets the session set by the nearest enclosing `runWithSession` call. + * @returns The current session, or `undefined` outside a transaction. + */ +export function getCurrentSession(): ClientSession | undefined { + return asyncLocalStorage.getStore(); +}