fix(instrumentation-mongoose): Fix instrumentation for Mongoose Model methods, save() and remove(), by passing options argument to original method calls

plugins/node/instrumentation-mongoose/src/mongoose.ts

@@ -330,6 +330,7 @@ export class MongooseInstrumentation extends InstrumentationBase {
           exec,
           originalThis,
           span,
+          args,
           self._config.responseHook,
           moduleVersion
         )

plugins/node/instrumentation-mongoose/src/utils.ts

@@ -98,16 +98,23 @@ export function handleCallbackResponse(
   exec: Function,
   originalThis: any,
   span: Span,
+  args: IArguments,
   responseHook?: MongooseResponseCustomAttributesFunction,
   moduleVersion: string | undefined = undefined
 ) {
-  return exec.apply(originalThis, [
-    (err: Error, response: any) => {
-      err
-        ? setErrorStatus(span, err)
-        : applyResponseHook(span, response, responseHook, moduleVersion);
-      span.end();
-      return callback!(err, response);
-    },
-  ]);
+  let callbackArgumentIndex = 0;
+  if (args.length === 2) {
+    callbackArgumentIndex = 1;
+  }
+
+  args[callbackArgumentIndex] = (err: Error, response: any): any => {
+    err
+      ? setErrorStatus(span, err)
+      : applyResponseHook(span, response, responseHook, moduleVersion);
+
+    span.end();
+    return callback!(err, response);
+  };
+
+  return exec.apply(originalThis, args);
 }

plugins/node/instrumentation-mongoose/test/mongoose.test.ts

@@ -66,8 +66,11 @@ describe('mongoose instrumentation', () => {
   beforeEach(async () => {
     instrumentation.disable();
     instrumentation.setConfig({
-      dbStatementSerializer: (_operation: string, payload) =>
-        JSON.stringify(payload),
+      dbStatementSerializer: (_operation: string, payload) => {
+        return JSON.stringify(payload, (key, value) => {
+          return key === 'session' ? '[Session]' : value;
+        });
+      },
     });
     instrumentation.enable();
     await loadUsers();
@@ -97,23 +100,159 @@ describe('mongoose instrumentation', () => {
     expect(statement.document).toEqual(expect.objectContaining(document));
   });
 
-  it('instrumenting save operation with callback', done => {
-    const document = {
-      firstName: 'Test first name',
-      lastName: 'Test last name',
-      email: 'test@example.com',
-    };
-    const user: IUser = new User(document);
+  describe('when save call does not have callback', async () => {
+    it('instrumenting save operation with promise and committed session', async () => {
+      const session = await User.startSession();
+      await session.startTransaction();
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: 'test@example.com',
+      };
+      const user: IUser = new User(document);
+
+      await user.save({ session });
+      await session.commitTransaction();
+      session.endSession();
 
-    user.save(() => {
       const spans = getTestSpans();
+      expect(spans.length).toBe(1);
+      assertSpan(spans[0] as ReadableSpan);
+      expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('save');
+      const statement = getStatement(spans[0] as ReadableSpan);
+      expect(statement.document).toEqual(expect.objectContaining(document));
+
+      const createdUser = await User.findById(user._id).lean();
+      expect(createdUser?._id.toString()).toEqual(user._id.toString());
+    });
 
+    it('instrumenting save operation with promise and session with aborted transaction ', async () => {
+      const session = await User.startSession();
+      await session.startTransaction();
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: 'test@example.com',
+      };
+      const user: IUser = new User(document);
+
+      await user.save({ session });
+      await session.abortTransaction();
+      session.endSession();
+
+      const spans = getTestSpans();
       expect(spans.length).toBe(1);
       assertSpan(spans[0] as ReadableSpan);
       expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('save');
       const statement = getStatement(spans[0] as ReadableSpan);
       expect(statement.document).toEqual(expect.objectContaining(document));
-      done();
+
+      const createdUser = await User.findById(user._id).lean();
+      expect(createdUser).toEqual(null);
+    });
+  });
+
+  describe('when save call has callback', async () => {
+    it('instrumenting save operation with promise, session and committed transaction', done => {
+      User.startSession().then(async session => {
+        await session.startTransaction();
+        const document = {
+          firstName: 'Test first name',
+          lastName: 'Test last name',
+          email: 'test@example.com',
+        };
+        const user: IUser = new User(document);
+        await user.save({ session }, async () => {
+          const spans = getTestSpans();
+          expect(spans.length).toBe(1);
+          assertSpan(spans[0] as ReadableSpan);
+          expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe(
+            'save'
+          );
+          const statement = getStatement(spans[0] as ReadableSpan);
+          expect(statement.document).toEqual(expect.objectContaining(document));
+
+          const createdUser = await User.findById(user._id).lean();
+          expect(createdUser?._id.toString()).toEqual(user._id.toString());
+          done();
+        });
+        await session.commitTransaction();
+        session.endSession();
+      });
+    });
+
+    it('instrumenting save operation with promise, session and aborted transaction', done => {
+      User.startSession().then(async session => {
+        await session.startTransaction();
+        const document = {
+          firstName: 'Test first name',
+          lastName: 'Test last name',
+          email: 'test@example.com',
+        };
+        const user: IUser = new User(document);
+
+        await user.save({ session }, async () => {
+          await session.abortTransaction();
+          session.endSession();
+
+          const spans = getTestSpans();
+          expect(spans.length).toBe(1);
+          assertSpan(spans[0] as ReadableSpan);
+          expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe(
+            'save'
+          );
+          const statement = getStatement(spans[0] as ReadableSpan);
+          expect(statement.document).toEqual(expect.objectContaining(document));
+
+          const createdUser = await User.findById(user._id).lean();
+          expect(createdUser).toEqual(null);
+          done();
+        });
+      });
+    });
+
+    it('instrumenting save operation with generic options and callback', done => {
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: 'test@example.com',
+      };
+      const user: IUser = new User(document);
+
+      user.save({}, () => {
+        const spans = getTestSpans();
+
+        expect(spans.length).toBe(1);
+        assertSpan(spans[0] as ReadableSpan);
+        expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe(
+          'save'
+        );
+        const statement = getStatement(spans[0] as ReadableSpan);
+        expect(statement.document).toEqual(expect.objectContaining(document));
+        done();
+      });
+    });
+
+    it('instrumenting save operation with only callback', done => {
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: 'test@example.com',
+      };
+      const user: IUser = new User(document);
+
+      user.save(() => {
+        const spans = getTestSpans();
+
+        expect(spans.length).toBe(1);
+        assertSpan(spans[0] as ReadableSpan);
+        expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe(
+          'save'
+        );
+        const statement = getStatement(spans[0] as ReadableSpan);
+        expect(statement.document).toEqual(expect.objectContaining(document));
+        done();
+      });
     });
   });