diff --git a/src/plugins/MongoosePlugin.ts b/src/plugins/MongoosePlugin.ts index 234abd5..e5348d4 100644 --- a/src/plugins/MongoosePlugin.ts +++ b/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(); diff --git a/tests/plugins/mongoose/client.ts b/tests/plugins/mongoose/client.ts index 25ff2b3..c3f4ed3 100644 --- a/tests/plugins/mongoose/client.ts +++ b/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', diff --git a/tests/plugins/mongoose/expected.data.yaml b/tests/plugins/mongoose/expected.data.yaml index 452d72b..a2400d8 100644 --- a/tests/plugins/mongoose/expected.data.yaml +++ b/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