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

* Fixing Mongoose instrumentation on Model methods with callback + tests

* Fixing and adding more tests

* Removing skip on existing failing test

* Renaming callback index variable

* use exported strings for attributes

* Defining args as IArguments and updating its values directly

* Fix lint issues

* Replacing tests with session option to use wtimeout to avoid the need of transactions in the instrumentation test

---------

Co-authored-by: Amir Blum <[email protected]>

Author: ferrelucas

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,90 @@ 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: '[email protected]',
-    };
-    const user: IUser = new User(document);
+  describe('when save call does not have callback', async () => {
+    it('instrumenting save operation with option property set', async () => {
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: '[email protected]',
+      };
+      const user: IUser = new User(document);
+      await user.save({ wtimeout: 42 });
 
-    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();
+      expect(statement.options.wtimeout).toEqual(42);
+
+      const createdUser = await User.findById(user._id).lean();
+      expect(createdUser?._id.toString()).toEqual(user._id.toString());
+    });
+  });
+
+  describe('when save call has callback', async () => {
+    it('instrumenting save operation with promise and option property set', done => {
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: '[email protected]',
+      };
+      const user: IUser = new User(document);
+      user.save({ wtimeout: 42 }, 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));
+        expect(statement.options.wtimeout).toEqual(42);
+
+        const createdUser = await User.findById(user._id).lean();
+        expect(createdUser?._id.toString()).toEqual(user._id.toString());
+        done();
+      });
+    });
+
+    it('instrumenting save operation with generic options and callback', done => {
+      const document = {
+        firstName: 'Test first name',
+        lastName: 'Test last name',
+        email: '[email protected]',
+      };
+      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: '[email protected]',
+      };
+      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();
+      });
     });
   });