Using mongodb transactions with typegoose

transaction.test.ts

import mongoose from "mongoose";
import { prop as Property, getModelForClass } from "@typegoose/typegoose";

import { withTransaction } from "./transaction";

class Foo {
  @Property({ required: true })
  example!: string;
}
const FooModel = getModelForClass(Foo);

class Bar {
  @Property({ required: true })
  example!: string;
}
const BarModel = getModelForClass(Bar);

describe("withTransaction", () => {
  beforeEach(async () => {
    await FooModel.createCollection();
    await BarModel.createCollection();
  });

  it("does not roll back changes when not used", async () => {
    const fn = async () => {
      await FooModel.create([{ example: "foo" }]);
      throw new Error();
    };

    await expect(fn).rejects.toThrow();

    expect(await FooModel.countDocuments({})).toBe(1);
  });

  describe("with single model", () => {
    it("rolls back changes when encountering error", async () => {
      const fn = async (session: mongoose.ClientSession) => {
        await FooModel.create([{ example: "foo" }], { session });
        throw new Error();
      };

      await expect(withTransaction(fn)).rejects.toThrow();

      expect(await FooModel.countDocuments({})).toBe(0);
    });

    it("resets mongoose documents when encountering error", async () => {
      const doc = await FooModel.create({ example: "foo" });

      const fn = async (session: mongoose.ClientSession) => {
        doc.example = "new";
        await doc.save({ session });
        throw new Error();
      };

      await expect(withTransaction(fn)).rejects.toThrow();

      expect(doc.modifiedPaths()).toEqual(["example"]);
    });

    it("commits changes without error", async () => {
      const fn = async (session: mongoose.ClientSession) => {
        await FooModel.create([{ example: "foo" }], { session });
      };

      await withTransaction(fn);

      expect(await FooModel.countDocuments({})).toBe(1);
    });
  });

  describe("with multiple models", () => {
    it("rolls back changes when encountering error", async () => {
      const fn = async (session: mongoose.ClientSession) => {
        await FooModel.create([{ example: "foo" }], { session });
        await BarModel.create([{ example: "bar" }], { session });
        throw new Error();
      };

      await expect(withTransaction(fn)).rejects.toThrow();

      expect(await FooModel.countDocuments({})).toBe(0);
      expect(await BarModel.countDocuments({})).toBe(0);
    });

    it("commits changes without error", async () => {
      const fn = async (session: mongoose.ClientSession) => {
        await FooModel.create([{ example: "foo" }], { session });
        await BarModel.create([{ example: "bar" }], { session });
      };

      await withTransaction(fn);

      expect(await FooModel.countDocuments({})).toBe(1);
      expect(await BarModel.countDocuments({})).toBe(1);
    });
  });

  describe("with nested transactions", () => {
    it("rolls back changes when encountering error", async () => {
      const fnFoo = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          await FooModel.create([{ example: "foo" }], { session });
        }, existingSession);
      };
      const fnBar = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          await BarModel.create([{ example: "bar" }], { session });
          throw new Error();
        }, existingSession);
      };
      const fn = async (session: mongoose.ClientSession) => {
        await fnFoo(session);
        await fnBar(session);
      };

      await expect(withTransaction(fn)).rejects.toThrow();

      expect(await FooModel.countDocuments({})).toBe(0);
      expect(await BarModel.countDocuments({})).toBe(0);
    });

    it("resets mongoose documents when encountering error", async () => {
      const foo = await FooModel.create({ example: "foo" });
      const bar = await BarModel.create({ example: "bar" });

      const fnFoo = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          foo.example = "new";
          await foo.save({ session });
        }, existingSession);
      };
      const fnBar = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          bar.example = "new";
          await bar.save({ session });
          throw new Error();
        }, existingSession);
      };
      const fn = async (session: mongoose.ClientSession) => {
        await fnFoo(session);
        await fnBar(session);
      };

      await expect(withTransaction(fn)).rejects.toThrow();

      expect(foo.modifiedPaths()).toEqual(["example"]);
      expect(bar.modifiedPaths()).toEqual(["example"]);
    });

    it("commits changes without error", async () => {
      const fnFoo = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          await FooModel.create([{ example: "foo" }], { session });
        }, existingSession);
      };
      const fnBar = async (existingSession?: mongoose.ClientSession) => {
        await withTransaction(async (session) => {
          await BarModel.create([{ example: "bar" }], { session });
        }, existingSession);
      };
      const fn = async (session: mongoose.ClientSession) => {
        await fnFoo(session);
        await fnBar(session);
      };

      await withTransaction(fn);

      expect(await FooModel.countDocuments({})).toBe(1);
      expect(await BarModel.countDocuments({})).toBe(1);
    });
  });
});

transaction.ts

import mongoose from "mongoose";

/*
 * See https://github.com/typegoose/typegoose/issues/279#issuecomment-645368737 and
 * https://thecodebarbarian.com/whats-new-in-mongoose-5-10-improved-transactions.html.
 */
export async function withTransaction(
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  fn: (session: mongoose.ClientSession) => Promise<any>,
  existingSession?: mongoose.ClientSession
): Promise<void> {
  if (existingSession) {
    if (existingSession.inTransaction()) return fn(existingSession);

    return mongoose.connection.transaction(fn);
  }

  const session = await mongoose.startSession();
  try {
    await mongoose.connection.transaction(fn);
  } finally {
    session.endSession();
  }
}