Skip to content

Commit 55cc256

Browse files
VovcharaaVovcharaa
authored
feat(instrumentation-mongoose): add instrumentation of static methods (#2748)
1 parent 11a2999 commit 55cc256

File tree

6 files changed

+241
-6
lines changed

6 files changed

+241
-6
lines changed

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,17 @@ export class MongooseInstrumentation extends InstrumentationBase<MongooseInstrum
153153
});
154154
this._wrap(moduleExports.Model, 'aggregate', this.patchModelAggregate());
155155

156+
this._wrap(
157+
moduleExports.Model,
158+
'insertMany',
159+
this.patchModelStatic('insertMany', moduleVersion)
160+
);
161+
this._wrap(
162+
moduleExports.Model,
163+
'bulkWrite',
164+
this.patchModelStatic('bulkWrite', moduleVersion)
165+
);
166+
156167
return moduleExports;
157168
}
158169

@@ -179,6 +190,9 @@ export class MongooseInstrumentation extends InstrumentationBase<MongooseInstrum
179190
this._unwrap(moduleExports.Query.prototype, funcName as any);
180191
});
181192
this._unwrap(moduleExports.Model, 'aggregate');
193+
194+
this._unwrap(moduleExports.Model, 'insertMany');
195+
this._unwrap(moduleExports.Model, 'bulkWrite');
182196
}
183197

184198
private patchAggregateExec(moduleVersion: string | undefined) {
@@ -314,6 +328,70 @@ export class MongooseInstrumentation extends InstrumentationBase<MongooseInstrum
314328
};
315329
}
316330

331+
private patchModelStatic(op: string, moduleVersion: string | undefined) {
332+
const self = this;
333+
return (original: Function) => {
334+
return function patchedStatic(
335+
this: any,
336+
docsOrOps: any,
337+
options?: any,
338+
callback?: Function
339+
) {
340+
if (
341+
self.getConfig().requireParentSpan &&
342+
trace.getSpan(context.active()) === undefined
343+
) {
344+
return original.apply(this, arguments);
345+
}
346+
if (typeof options === 'function') {
347+
callback = options;
348+
options = undefined;
349+
}
350+
351+
const serializePayload: SerializerPayload = {};
352+
switch (op) {
353+
case 'insertMany':
354+
serializePayload.documents = docsOrOps;
355+
break;
356+
case 'bulkWrite':
357+
serializePayload.operations = docsOrOps;
358+
break;
359+
default:
360+
serializePayload.document = docsOrOps;
361+
break;
362+
}
363+
if (options !== undefined) {
364+
serializePayload.options = options;
365+
}
366+
367+
const attributes: Attributes = {};
368+
const { dbStatementSerializer } = self.getConfig();
369+
if (dbStatementSerializer) {
370+
attributes[SEMATTRS_DB_STATEMENT] = dbStatementSerializer(
371+
op,
372+
serializePayload
373+
);
374+
}
375+
376+
const span = self._startSpan(
377+
this.collection,
378+
this.modelName,
379+
op,
380+
attributes
381+
);
382+
383+
return self._handleResponse(
384+
span,
385+
original,
386+
this,
387+
arguments,
388+
callback,
389+
moduleVersion
390+
);
391+
};
392+
};
393+
}
394+
317395
// we want to capture the otel span on the object which is calling exec.
318396
// in the special case of aggregate, we need have no function to path
319397
// on the Aggregate object to capture the context on, so we patch

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ export interface SerializerPayload {
2323
document?: any;
2424
aggregatePipeline?: any;
2525
fields?: any;
26+
documents?: any;
27+
operations?: any;
2628
}
2729

2830
export type DbStatementSerializer = (

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,8 @@ export function handleCallbackResponse(
105105
let callbackArgumentIndex = 0;
106106
if (args.length === 2) {
107107
callbackArgumentIndex = 1;
108+
} else if (args.length === 3) {
109+
callbackArgumentIndex = 2;
108110
}
109111

110112
args[callbackArgumentIndex] = (err: Error, response: any): any => {

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

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,8 @@ describe('mongoose instrumentation [common]', () => {
7474
});
7575
},
7676
});
77-
instrumentation.enable();
7877
await loadUsers();
79-
await User.createIndexes();
78+
instrumentation.enable();
8079
});
8180

8281
afterEach(async () => {
@@ -321,6 +320,107 @@ describe('mongoose instrumentation [common]', () => {
321320
expect(statement.document).toEqual(expect.objectContaining(document));
322321
});
323322

323+
it('instrumenting insertMany operation', async () => {
324+
const documents = [
325+
{
326+
firstName: 'John',
327+
lastName: 'Doe',
328+
email: '[email protected]',
329+
},
330+
{
331+
firstName: 'Jane',
332+
lastName: 'Doe',
333+
email: '[email protected]',
334+
},
335+
];
336+
await User.insertMany(documents);
337+
338+
const spans = getTestSpans();
339+
expect(spans.length).toBe(1);
340+
assertSpan(spans[0] as ReadableSpan);
341+
expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('insertMany');
342+
const statement = getStatement(spans[0] as ReadableSpan);
343+
expect(statement.documents).toEqual(documents);
344+
});
345+
346+
it('instrumenting bulkWrite operation', async () => {
347+
const operations = [
348+
{
349+
insertOne: {
350+
document: {
351+
firstName: 'Jane',
352+
lastName: 'Doe',
353+
email: '[email protected]',
354+
age: 25,
355+
},
356+
},
357+
},
358+
{
359+
updateMany: {
360+
filter: { age: { $lte: 20 } },
361+
update: { $set: { age: 20 } },
362+
},
363+
},
364+
{
365+
updateOne: {
366+
filter: { firstName: 'Jane' },
367+
update: { $inc: { age: 1 } },
368+
},
369+
},
370+
{ deleteOne: { filter: { firstName: 'Michael' } } },
371+
{
372+
updateOne: {
373+
filter: { firstName: 'Zara' },
374+
update: {
375+
$set: { lastName: 'Doe', age: 40, email: '[email protected]' },
376+
},
377+
upsert: true,
378+
},
379+
},
380+
];
381+
await User.bulkWrite(operations);
382+
383+
const spans = getTestSpans();
384+
expect(spans.length).toBe(1);
385+
assertSpan(spans[0] as ReadableSpan);
386+
expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('bulkWrite');
387+
const statement = getStatement(spans[0] as ReadableSpan);
388+
expect(statement.operations).toEqual([
389+
{
390+
insertOne: {
391+
document: {
392+
firstName: 'Jane',
393+
lastName: 'Doe',
394+
email: '[email protected]',
395+
age: 25,
396+
},
397+
},
398+
},
399+
{
400+
updateMany: {
401+
filter: { age: { $lte: 20 } },
402+
update: { $set: { age: 20 } },
403+
},
404+
},
405+
{
406+
updateOne: {
407+
filter: { firstName: 'Jane' },
408+
update: { $inc: { age: 1 } },
409+
},
410+
},
411+
{ deleteOne: { filter: { firstName: 'Michael' } } },
412+
{
413+
updateOne: {
414+
filter: { firstName: 'Zara' },
415+
update: {
416+
$set: { lastName: 'Doe', age: 40, email: '[email protected]' },
417+
},
418+
upsert: true,
419+
},
420+
},
421+
]);
422+
});
423+
324424
it('instrumenting aggregate operation', async () => {
325425
await User.aggregate([
326426
{ $match: { firstName: 'John' } },

‎plugins/node/instrumentation-mongoose/test/mongoose-v5-v6.test.ts

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,8 @@ describe('mongoose instrumentation [v5/v6]', () => {
7171
});
7272
},
7373
});
74-
instrumentation.enable();
7574
await loadUsers();
76-
await User.createIndexes();
75+
instrumentation.enable();
7776
});
7877

7978
afterEach(async () => {
@@ -152,6 +151,61 @@ describe('mongoose instrumentation [v5/v6]', () => {
152151
});
153152
});
154153

154+
describe('when insertMany call has callback', async () => {
155+
it('instrumenting insertMany operation with generic options and callback', done => {
156+
const documents = [
157+
{
158+
firstName: 'John',
159+
lastName: 'Doe',
160+
email: '[email protected]',
161+
},
162+
{
163+
firstName: 'Jane',
164+
lastName: 'Doe',
165+
email: '[email protected]',
166+
},
167+
];
168+
// @ts-ignore - v7 removed callback support
169+
// https://mongoosejs.com/docs/migrating_to_7.html#dropped-callback-support
170+
User.insertMany(documents, { ordered: true }, () => {
171+
const spans = getTestSpans();
172+
expect(spans.length).toBe(1);
173+
assertSpan(spans[0] as ReadableSpan);
174+
expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('insertMany');
175+
const statement = getStatement(spans[0] as ReadableSpan);
176+
expect(statement.documents).toEqual(documents);
177+
expect(statement.options.ordered).toEqual(true);
178+
done();
179+
});
180+
});
181+
182+
it('instrumenting insertMany operation with only callback', done => {
183+
const documents = [
184+
{
185+
firstName: 'John',
186+
lastName: 'Doe',
187+
email: '[email protected]',
188+
},
189+
{
190+
firstName: 'Jane',
191+
lastName: 'Doe',
192+
email: '[email protected]',
193+
},
194+
];
195+
// @ts-ignore - v7 removed callback support
196+
// https://mongoosejs.com/docs/migrating_to_7.html#dropped-callback-support
197+
User.insertMany(documents, () => {
198+
const spans = getTestSpans();
199+
expect(spans.length).toBe(1);
200+
assertSpan(spans[0] as ReadableSpan);
201+
expect(spans[0].attributes[SEMATTRS_DB_OPERATION]).toBe('insertMany');
202+
const statement = getStatement(spans[0] as ReadableSpan);
203+
expect(statement.documents).toEqual(documents);
204+
done();
205+
});
206+
});
207+
});
208+
155209
describe('remove operation', () => {
156210
it('instrumenting remove operation [deprecated]', async () => {
157211
const user = await User.findOne({ email: '[email protected]' });

‎plugins/node/instrumentation-mongoose/test/mongoose-v7-v8.test.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,8 @@ describe('mongoose instrumentation [v7/v8]', () => {
6868
});
6969
},
7070
});
71-
instrumentation.enable();
7271
await loadUsers();
73-
await User.createIndexes();
72+
instrumentation.enable();
7473
});
7574

7675
afterEach(async () => {

0 commit comments

Comments
 (0)