Skip to content

Commit 1ece354

Browse files
committed
feat: MongoDB Tracing Support
1 parent 7dd42bc commit 1ece354

File tree

4 files changed

+315
-1
lines changed

4 files changed

+315
-1
lines changed

packages/tracing/src/integrations/express.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ export class Express implements Integration {
8585
*/
8686
public setupOnce(): void {
8787
if (!this._app) {
88-
logger.error('ExpressIntegration is missing an Express instance');
88+
logger.error('Express Integration is missing an Express instance');
8989
return;
9090
}
9191
instrumentMiddlewares(this._app);
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
export { Express } from './express';
2+
export { Mongo } from './mongo';
Lines changed: 232 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,232 @@
1+
import { Hub } from '@sentry/hub';
2+
import { EventProcessor, Integration, SpanContext } from '@sentry/types';
3+
import { fill, logger } from '@sentry/utils';
4+
5+
// TODO: Support Cursors? — Kamil
6+
type Operation =
7+
| 'aggregate' // aggregate(pipeline, options, callback)
8+
| 'bulkWrite' // bulkWrite(operations, options, callback)
9+
| 'countDocuments' // countDocuments(query, options, callback)
10+
| 'createIndex' // createIndex(fieldOrSpec, options, callback)
11+
| 'createIndexes' // createIndexes(indexSpecs, options, callback)
12+
| 'deleteMany' // deleteMany(filter, options, callback)
13+
| 'deleteOne' // deleteOne(filter, options, callback)
14+
| 'distinct' // distinct(key, query, options, callback)
15+
| 'drop' // drop(options, callback)
16+
| 'dropIndex' // dropIndex(indexName, options, callback)
17+
| 'dropIndexes' // dropIndexes(options, callback)
18+
| 'estimatedDocumentCount' // estimatedDocumentCount(options, callback)
19+
| 'findOne' // findOne(query, options, callback)
20+
| 'findOneAndDelete' // findOneAndDelete(filter, options, callback)
21+
| 'findOneAndReplace' // findOneAndReplace(filter, replacement, options, callback)
22+
| 'findOneAndUpdate' // findOneAndUpdate(filter, update, options, callback)
23+
| 'indexes' // indexes(options, callback)
24+
| 'indexExists' // indexExists(indexes, options, callback)
25+
| 'indexInformation' // indexInformation(options, callback)
26+
| 'initializeOrderedBulkOp' // initializeOrderedBulkOp(options, callback)
27+
| 'insertMany' // insertMany(docs, options, callback)
28+
| 'insertOne' // insertOne(doc, options, callback)
29+
| 'isCapped' // isCapped(options, callback)
30+
| 'mapReduce' // mapReduce(map, reduce, options, callback)
31+
| 'options' // options(options, callback)
32+
| 'parallelCollectionScan' // parallelCollectionScan(options, callback)
33+
| 'rename' // rename(newName, options, callback)
34+
| 'replaceOne' // replaceOne(filter, doc, options, callback)
35+
| 'stats' // stats(options, callback)
36+
| 'updateMany' // updateMany(filter, update, options, callback)
37+
| 'updateOne'; // updateOne(filter, update, options, callback)
38+
39+
// 2 arguments operations which accept no input which can be used to describe the operation.
40+
const simpleOperations: Operation[] = [
41+
'drop',
42+
'dropIndexes',
43+
'estimatedDocumentCount',
44+
'indexes',
45+
'indexInformation',
46+
'initializeOrderedBulkOp',
47+
'isCapped',
48+
'options',
49+
'parallelCollectionScan',
50+
'stats',
51+
];
52+
53+
// 4 arguments operations which accept 2 inputs which can be used to describe the operation.
54+
const complexOperations: Operation[] = [
55+
'distinct',
56+
'findOneAndReplace',
57+
'findOneAndUpdate',
58+
'replaceOne',
59+
'updateMany',
60+
'updateOne',
61+
];
62+
63+
const operationSignatures: {
64+
[op in Operation]?: string[];
65+
} = {
66+
bulkWrite: ['operations'],
67+
countDocuments: ['query'],
68+
createIndex: ['fieldOrSpec'],
69+
createIndexes: ['indexSpecs'],
70+
deleteMany: ['filter'],
71+
deleteOne: ['filter'],
72+
distinct: ['key', 'query'],
73+
dropIndex: ['indexName'],
74+
findOne: ['query'],
75+
findOneAndDelete: ['filter'],
76+
findOneAndReplace: ['filter', 'replacement'],
77+
findOneAndUpdate: ['filter', 'update'],
78+
indexExists: ['indexes'],
79+
insertMany: ['docs'],
80+
insertOne: ['doc'],
81+
mapReduce: ['map', 'reduce'],
82+
rename: ['newName'],
83+
replaceOne: ['filter', 'doc'],
84+
updateMany: ['filter', 'update'],
85+
updateOne: ['filter', 'update'],
86+
};
87+
88+
interface Collection {
89+
collectionName: string;
90+
dbName: string;
91+
namespace: string;
92+
prototype: {
93+
[operation in Operation]: (...args: unknown[]) => unknown;
94+
};
95+
}
96+
97+
interface MongoOptions {
98+
collection?: Collection;
99+
operations?: Operation[];
100+
describeOperations?: boolean | Operation[];
101+
}
102+
103+
/** Tracing integration for node-postgres package */
104+
export class Mongo implements Integration {
105+
/**
106+
* @inheritDoc
107+
*/
108+
public static id: string = 'Mongo';
109+
110+
/**
111+
* @inheritDoc
112+
*/
113+
public name: string = Mongo.id;
114+
115+
private _collection?: Collection;
116+
private _operations: Operation[];
117+
private _describeOperations?: boolean | Operation[];
118+
119+
/**
120+
* @inheritDoc
121+
*/
122+
public constructor(options: MongoOptions = {}) {
123+
this._collection = options.collection;
124+
this._operations = options.operations || ['findOne'];
125+
this._describeOperations = 'describeOperations' in options ? options.describeOperations : true;
126+
}
127+
128+
/**
129+
* @inheritDoc
130+
*/
131+
public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
132+
if (!this._collection) {
133+
logger.error('Mongo Integration is missing a Mongo.Collection constructor');
134+
return;
135+
}
136+
this._instrumentOperations(this._collection, this._operations, getCurrentHub);
137+
}
138+
139+
/**
140+
* Patches original collection methods
141+
*/
142+
private _instrumentOperations(collection: Collection, operations: Operation[], getCurrentHub: () => Hub): void {
143+
operations.forEach((operation: Operation) => this._patchOperation(collection, operation, getCurrentHub));
144+
}
145+
146+
/**
147+
* Patches original collection to utilize our tracing functionality
148+
*/
149+
private _patchOperation(collection: Collection, operation: Operation, getCurrentHub: () => Hub): void {
150+
if (!(operation in collection.prototype)) return;
151+
152+
const getSpanContext = this._getSpanContextFromOperationArguments.bind(this);
153+
154+
fill(collection.prototype, operation, function(orig: () => void | Promise<unknown>) {
155+
return function(this: Collection, ...args: unknown[]) {
156+
const lastArg = args[args.length - 1];
157+
const scope = getCurrentHub().getScope();
158+
const transaction = scope?.getTransaction();
159+
160+
// mapReduce is a special edge-case, as it's the only operation that accepts functions
161+
// other than the callback as it's own arguments. Therefore despite lastArg being
162+
// a function, it can be still a promise-based call without a callback.
163+
// mapReduce(map, reduce, options, callback) where `[map|reduce]: function | string`
164+
if (typeof lastArg !== 'function' || (operation === 'mapReduce' && args.length === 2)) {
165+
const span = transaction?.startChild(getSpanContext(this, operation, args));
166+
return (orig.call(this, ...args) as Promise<unknown>).then((res: unknown) => {
167+
span?.finish();
168+
return res;
169+
});
170+
}
171+
172+
const span = transaction?.startChild(getSpanContext(this, operation, args.slice(0, -1)));
173+
return orig.call(this, ...args.slice(0, -1), function(err: Error, result: unknown) {
174+
span?.finish();
175+
lastArg(err, result);
176+
});
177+
};
178+
});
179+
}
180+
181+
/**
182+
* Form a SpanContext based on the user input to a given operation.
183+
*/
184+
private _getSpanContextFromOperationArguments(
185+
collection: Collection,
186+
operation: Operation,
187+
args: unknown[],
188+
): SpanContext {
189+
const data: { [key: string]: string } = {
190+
collectionName: collection.collectionName,
191+
dbName: collection.dbName,
192+
namespace: collection.namespace,
193+
};
194+
const spanContext: SpanContext = {
195+
op: `query.${operation}`,
196+
data,
197+
};
198+
199+
// There's nothing we can extract from these operations, other than the name
200+
// of the collection it was invoked on, so just return early.
201+
// Or there's no signature available for us to be used for the extracted data description.
202+
const signature = operationSignatures[operation];
203+
if (simpleOperations.includes(operation) || !signature) {
204+
return spanContext;
205+
}
206+
207+
const shouldDescribe = Array.isArray(this._describeOperations)
208+
? this._describeOperations.includes(operation)
209+
: this._describeOperations;
210+
211+
// Special case for `mapReduce`, as the only one accepting functions as arguments.
212+
if (operation === 'mapReduce' && shouldDescribe) {
213+
const [map, reduce] = args as { name?: string }[];
214+
data[signature[0]] = typeof map === 'string' ? map : map.name || '<anonymous>';
215+
data[signature[1]] = typeof reduce === 'string' ? reduce : reduce.name || '<anonymous>';
216+
return spanContext;
217+
}
218+
219+
if (shouldDescribe) {
220+
try {
221+
data[signature[0]] = JSON.stringify(args[0]);
222+
if (complexOperations.includes(operation)) {
223+
data[signature[1]] = JSON.stringify(args[1]);
224+
}
225+
} catch (o_O) {
226+
// no-empty
227+
}
228+
}
229+
230+
return spanContext;
231+
}
232+
}
Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import { Hub } from '@sentry/hub';
2+
import { EventProcessor, Integration } from '@sentry/types';
3+
import { fill, logger } from '@sentry/utils';
4+
5+
interface PgClient {
6+
prototype: {
7+
query: () => void | Promise<unknown>;
8+
};
9+
}
10+
11+
interface PostgresOptions {
12+
client?: PgClient;
13+
}
14+
15+
/** Tracing integration for node-postgres package */
16+
export class Postgres implements Integration {
17+
/**
18+
* @inheritDoc
19+
*/
20+
public static id: string = 'Postgres';
21+
22+
/**
23+
* @inheritDoc
24+
*/
25+
public name: string = Postgres.id;
26+
27+
private _client?: PgClient;
28+
29+
/**
30+
* @inheritDoc
31+
*/
32+
public constructor(options: PostgresOptions = {}) {
33+
this._client = options.client;
34+
}
35+
36+
/**
37+
* @inheritDoc
38+
*/
39+
public setupOnce(_: (callback: EventProcessor) => void, getCurrentHub: () => Hub): void {
40+
if (!this._client) {
41+
logger.error('Postgres Integration is missing a Postgres.Client constructor');
42+
return;
43+
}
44+
45+
/**
46+
* function (query, callback) => void
47+
* function (query, params, callback) => void
48+
* function (query) => Promise
49+
* function (query, params) => Promise
50+
*/
51+
fill(this._client.prototype, 'query', function(orig: () => void | Promise<unknown>) {
52+
return function(this: unknown, config: unknown, values: unknown, callback: unknown) {
53+
const scope = getCurrentHub().getScope();
54+
const transaction = scope?.getTransaction();
55+
const span = transaction?.startChild({
56+
description: typeof config === 'string' ? config : (config as { text: string }).text,
57+
op: `query`,
58+
});
59+
60+
if (typeof callback === 'function') {
61+
return orig.call(this, config, values, function(err: Error, result: unknown) {
62+
if (span) span.finish();
63+
callback(err, result);
64+
});
65+
}
66+
67+
if (typeof values === 'function') {
68+
return orig.call(this, config, function(err: Error, result: unknown) {
69+
if (span) span.finish();
70+
values(err, result);
71+
});
72+
}
73+
74+
return (orig.call(this, config, values) as Promise<unknown>).then((res: unknown) => {
75+
if (span) span.finish();
76+
return res;
77+
});
78+
};
79+
});
80+
}
81+
}

0 commit comments

Comments
 (0)