Modify mongoose call xxxx() throwing "is not a function" (#128)

Author: shivendoodeshmukh

src/plugins/MongoosePlugin.ts

@@ -118,7 +118,6 @@ class MongoosePlugin implements SwPlugin {
           arguments[arguments.length - 1] = function () {
             // in case of immediate synchronous callback from mongoose
             (span as any).mongooseInCall = false;
-
             wrappedCallback.apply(this, arguments as any);
           };
         }
@@ -130,7 +129,42 @@ class MongoosePlugin implements SwPlugin {
         if (!hasCB) {
           if (ret && typeof ret.then === 'function') {
             // generic Promise check
-            ret = wrapPromise(span, ret);
+
+            if (ret.constructor.name != 'Query') {
+              ret = wrapPromise(span, ret);
+            } else {
+              // Mongoose Query object
+              const chainMethods = ['select', 'sort', 'skip', 'limit', 'populate'];
+
+              // Mongoose Query object
+              const originalThen = ret.then;
+              const originalExec = ret.exec;
+              const originalLean = ret.lean;
+
+              // Preserve the query chain methods using arrow functions to maintain context
+              ret.then = (...args: any[]) => wrapPromise(span, originalThen.apply(ret, args));
+              ret.exec = (...args: any[]) => wrapPromise(span, originalExec.apply(ret, args));
+              ret.lean = (...args: any[]) => {
+                const leanQuery = originalLean.apply(ret, args);
+                // Preserve other chain methods on the lean result
+                leanQuery.then = ret.then;
+                leanQuery.exec = ret.exec;
+                return leanQuery;
+              };
+              // Wrap other common query methods that might be chained
+              chainMethods.forEach((method) => {
+                if (ret[method]) {
+                  const originalMethod = ret[method];
+                  ret[method] = (...args: any[]) => {
+                    const result = originalMethod.apply(ret, args);
+                    result.then = ret.then;
+                    result.exec = ret.exec;
+                    result.lean = ret.lean;
+                    return result;
+                  };
+                }
+              });
+            }
           } else {
             // no callback passed in and no Promise or Cursor returned, play it safe
             span.stop();

tests/plugins/mongoose/client.ts

@@ -21,6 +21,7 @@ import * as http from 'http';
 import agent from '../../../src';
 
 process.env.SW_AGENT_LOGGING_LEVEL = 'ERROR';
+process.env.SW_AGENT_DISABLE_PLUGINS = 'MongoDBPlugin';
 
 agent.start({
   serviceName: 'client',

tests/plugins/mongoose/expected.data.yaml

@@ -35,7 +35,7 @@ segmentItems:
             tags:
               - { key: db.type, value: MongoDB }
               - { key: db.instance, value: admin }
-              - { key: db.statement, value: 'collection("tests")' }
+              - { key: db.statement, 'collection("tests")' }
           - operationName: Mongoose/ensureIndexes
             operationId: 0
             parentSpanId: 0
@@ -50,6 +50,21 @@ segmentItems:
             tags:
               - { key: db.type, value: MongoDB }
               - { key: db.instance, value: admin }
+          - operationName: MongoDB/find         # This is the required span
+            operationId: 0
+            parentSpanId: 0
+            spanId: 4
+            spanLayer: Database
+            startTime: gt 0
+            endTime: gt 0
+            componentId: 9
+            spanType: Exit
+            peer: mongo:27017
+            skipAnalysis: false
+            tags:
+              - { key: db.type, value: MongoDB }
+              - { key: db.instance, value: admin }
+              - { key: db.statement, 'tests.find({})' }
           - operationName: Mongoose/find
             operationId: 0
             parentSpanId: 0
@@ -64,7 +79,7 @@ segmentItems:
             tags:
               - { key: db.type, value: MongoDB }
               - { key: db.instance, value: admin }
-              - { key: db.statement, value: 'tests.find({})' }
+              # The mongoose find span might not have the db.statement, which is fine
           - operationName: GET:/mongoose
             operationId: 0
             parentSpanId: -1