Skip to content

Commit

Permalink
feat: add mongodb plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
vmarchaud committed Aug 21, 2019
1 parent 7f5d71b commit 0af967e
Show file tree
Hide file tree
Showing 2 changed files with 260 additions and 2 deletions.
7 changes: 5 additions & 2 deletions packages/opentelemetry-plugin-mongodb/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@
},
"devDependencies": {
"@types/mocha": "^5.2.7",
"@types/node": "^12.6.9",
"@types/mongodb": "^3.2.3",
"@types/node": "^12.7.2",
"@types/shimmer": "^1.0.1",
"codecov": "^3.5.0",
"gts": "^1.1.0",
"mocha": "^6.2.0",
Expand All @@ -51,6 +53,7 @@
"dependencies": {
"@opentelemetry/core": "^0.0.1",
"@opentelemetry/node-tracer": "^0.0.1",
"@opentelemetry/types": "^0.0.1"
"@opentelemetry/types": "^0.0.1",
"shimmer": "^1.2.1"
}
}
255 changes: 255 additions & 0 deletions packages/opentelemetry-plugin-mongodb/src/mongodb.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,255 @@
/**
* Copyright 2018, OpenCensus Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

// mongodb.Server type is deprecated so every use trigger a lint error
/* tslint:disable:deprecation */

import { BasePlugin, NoopLogger } from '@opentelemetry/core';
import {
Span,
SpanKind,
Logger,
Tracer,
CanonicalCode,
} from '@opentelemetry/types';
import * as mongodb from 'mongodb';
import * as shimmer from 'shimmer';

type Func<T> = (...args: unknown[]) => T;
type MongoDB = typeof mongodb;
interface MongoInternalCommand {
findandmodify: boolean;
createIndexes: boolean;
count: boolean;
ismaster: boolean;
}

/** MongoDB instrumentation plugin for Opencensus */
export class MongoDBPlugin extends BasePlugin<MongoDB> {
private readonly SERVER_FNS = ['insert', 'update', 'remove', 'auth'];
private readonly CURSOR_FNS_FIRST = ['_find', '_getmore'];
protected _logger!: Logger;
protected readonly _tracer!: Tracer;

constructor(public moduleName: string) {
super();
// TODO: remove this once a logger will be passed
// https://github.com/open-telemetry/opentelemetry-js/issues/193
this._logger = new NoopLogger();
}

/**
* Patches MongoDB operations.
*/
protected patch() {
this._logger.debug('Patched MongoDB');

if (this._moduleExports.Server) {
this._logger.debug('patching mongodb-core.Server.prototype.command');
shimmer.wrap(
this._moduleExports.Server.prototype,
'command' as never,
// tslint:disable-next-line:no-any
this.getPatchCommand() as any
);
this._logger.debug(
'patching mongodb-core.Server.prototype functions:',
this.SERVER_FNS
);
shimmer.massWrap(
[this._moduleExports.Server.prototype],
this.SERVER_FNS as never[],
// tslint:disable-next-line:no-any
this.getPatchQuery() as any
);
}

if (this._moduleExports.Cursor) {
this._logger.debug(
'patching mongodb-core.Cursor.prototype functions:',
this.CURSOR_FNS_FIRST
);
shimmer.massWrap(
[this._moduleExports.Cursor.prototype],
this.CURSOR_FNS_FIRST as never[],
// tslint:disable-next-line:no-any
this.getPatchCursor() as any
);
}

return this._moduleExports;
}

/** Unpatches all MongoDB patched functions. */
unpatch(): void {
shimmer.unwrap(this._moduleExports.Server.prototype, 'command' as never);
shimmer.massUnwrap([this._moduleExports.Server.prototype], this
.SERVER_FNS as never[]);
shimmer.massUnwrap([this._moduleExports.Cursor.prototype], this
.CURSOR_FNS_FIRST as never[]);
}

/** Creates spans for Command operations */
private getPatchCommand() {
const plugin = this;
return (original: Func<mongodb.Server>) => {
return function(
this: mongodb.Server,
ns: string,
command: MongoInternalCommand,
options: {} | Function,
callback: Func<unknown>
): mongodb.Server {
const currentSpan = plugin._tracer.getCurrentSpan();
if (currentSpan === null) {
return original.apply(this, (arguments as unknown) as unknown[]);
}
const resultHandler =
typeof options === 'function' ? options : callback;
if (typeof resultHandler !== 'function') {
return original.apply(this, (arguments as unknown) as unknown[]);
}
if (typeof command !== 'object') {
return original.apply(this, (arguments as unknown) as unknown[]);
}
let type: string;
if (command.createIndexes) {
type = 'createIndexes';
} else if (command.findandmodify) {
type = 'findAndModify';
} else if (command.ismaster) {
type = 'isMaster';
} else if (command.count) {
type = 'count';
} else {
type = 'command';
}

const span = plugin._tracer.startSpan(`${ns}.${type}`, {
parent: currentSpan,
kind: SpanKind.CLIENT,
});
if (typeof options === 'function') {
return original.call(
this,
ns,
command,
plugin.patchEnd(span, options as Func<unknown>)
);
} else {
return original.call(
this,
ns,
command,
options,
plugin.patchEnd(span, callback)
);
}
};
};
}

/** Creates spans for Query operations */
private getPatchQuery() {
const plugin = this;
return (original: Func<mongodb.Server>) => {
return function(
this: mongodb.Server,
ns: string,
command: MongoInternalCommand,
options: {},
callback: Func<unknown>
): mongodb.Server {
const currentSpan = plugin._tracer.getCurrentSpan();
if (currentSpan === null) {
return original.apply(this, (arguments as unknown) as unknown[]);
}
const resultHandler =
typeof options === 'function' ? options : callback;
if (typeof resultHandler !== 'function') {
return original.apply(this, (arguments as unknown) as unknown[]);
}
const span = plugin._tracer.startSpan(`${ns}.query`, {
kind: SpanKind.CLIENT,
parent: currentSpan,
});
if (typeof options === 'function') {
return original.call(
this,
ns,
command,
plugin.patchEnd(span, options as Func<unknown>)
);
} else {
return original.call(
this,
ns,
command,
options,
plugin.patchEnd(span, callback)
);
}
};
};
}

/** Creates spans for Cursor operations */
private getPatchCursor() {
const plugin = this;
return (original: Func<mongodb.Cursor>) => {
return function(
this: { ns: string },
...args: unknown[]
): mongodb.Cursor {
const currentSpan = plugin._tracer.getCurrentSpan();
if (currentSpan === null) {
return original.apply(this, (arguments as unknown) as unknown[]);
}
const resultHandler = args[0] as Func<unknown> | undefined;
if (resultHandler === undefined) {
return original.apply(this, (arguments as unknown) as unknown[]);
}
const span = plugin._tracer.startSpan(`${this.ns}.cursor`, {
parent: currentSpan,
kind: SpanKind.CLIENT,
});
return original.call(this, plugin.patchEnd(span, resultHandler));
};
};
}

/**
* Ends a created span.
* @param span The created span to end.
* @param resultHandler A callback function.
*/
patchEnd(span: Span, resultHandler: Func<unknown>): Function {
return function patchedEnd(this: {}, ...args: unknown[]) {
const error = args[0];
if (error instanceof Error) {
span.setStatus({
code: CanonicalCode.UNKNOWN,
message: error.message,
});
}
span.end();
return resultHandler.apply(this, (arguments as unknown) as unknown[]);
};
}
}

const plugin = new MongoDBPlugin('mongodb');
export { plugin };

0 comments on commit 0af967e

Please sign in to comment.