feat: add transactional support

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,

src/common/decorators/transactional.decorator.spec.ts

@@ -0,0 +1,82 @@
+import { Transactional } from "./transactional.decorator";
+import { MongoTransactionModule } from "../modules/mongo-transaction.module";
+import { MongoTransactionService } from "../services/mongo-transaction.service";
+import { getCurrentSession } from "../utils/session-context.util";
+
+function registerMongoTransactionService(service: MongoTransactionService) {
+  new MongoTransactionModule(service).onModuleInit();
+}
+
+describe("Transactional", () => {
+  class TestService {
+    @Transactional()
+    async withArgs(a: string, b: string) {
+      return { a, b };
+    }
+  }
+
+  let service: TestService;
+  const fakeMongoTransactionService = {
+    run: jest.fn().mockImplementation((fn) => fn()),
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    fakeMongoTransactionService.run.mockImplementation((fn) => fn());
+    registerMongoTransactionService(fakeMongoTransactionService as never);
+    service = new TestService();
+  });
+
+  it("delegates to MongoTransactionService.run and forwards all arguments unchanged", async () => {
+    const result = await service.withArgs("x", "y");
+
+    expect(fakeMongoTransactionService.run).toHaveBeenCalledTimes(1);
+    expect(result).toEqual({ a: "x", b: "y" });
+  });
+
+  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 {
+      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);
+      registerMongoTransactionService(
+        new MongoTransactionService(connection as never),
+      );
+    });
+
+    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]);
+    });
+  });
+});

src/common/decorators/transactional.decorator.ts

@@ -0,0 +1,28 @@
+import { getMongoTransactionService } from "../modules/mongo-transaction.module";
+
+/**
+ * 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.
+ */
+export function Transactional(): MethodDecorator {
+  return function (
+    _target: object,
+    _propertyKey: string | symbol,
+    descriptor: PropertyDescriptor,
+  ) {
+    const originalMethod = descriptor.value;
+
+    descriptor.value = function (this: unknown, ...args: unknown[]) {
+      return getMongoTransactionService().run(() =>
+        originalMethod.apply(this, args),
+      );
+    };
+
+    return descriptor;
+  };
+}

src/common/modules/mongo-transaction.module.ts

@@ -0,0 +1,35 @@
+import { Global, Module, OnModuleInit } from "@nestjs/common";
+import { MongoTransactionService } from "../services/mongo-transaction.service";
+
+let instance: MongoTransactionService | undefined;
+
+@Global()
+@Module({
+  providers: [MongoTransactionService],
+  exports: [MongoTransactionService],
+})
+export class MongoTransactionModule implements OnModuleInit {
+  constructor(
+    private readonly mongoTransactionService: MongoTransactionService,
+  ) {}
+
+  onModuleInit() {
+    instance = this.mongoTransactionService;
+  }
+}
+
+/**
+ * Returns the MongoTransactionService singleton registered by
+ * MongoTransactionModule. Backs `@Transactional()` so decorated methods
+ * don't rely on the owning class injecting MongoTransactionService itself.
+ */
+export function getMongoTransactionService(): MongoTransactionService {
+  if (!instance) {
+    throw new Error(
+      "MongoTransactionService is not available yet — make sure " +
+        "MongoTransactionModule is imported in AppModule before any " +
+        "@Transactional() method runs.",
+    );
+  }
+  return instance;
+}

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