/**
 * 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.
 */ import { BasePlugin , Func , Span , SpanKind } from '@opencensus/core' ; import * as shimmer from 'shimmer' ; export type MongoPluginConfig = { /**
     * Add arguments to the span metadata for a every command
     */ detailedCommands : boolean ; } ; /** MongoDB instrumentation plugin for Opencensus */ export class MongoDBPlugin extends BasePlugin { protected options : MongoPluginConfig ; protected readonly internalFileList = { '1 - 3' : { 'ConnectionPool' : 'lib/connection/pool' } } ; /** Constructs a new MongoDBPlugin instance. */ constructor ( moduleName ) { super ( moduleName ) ; } /**
     * Patches MongoDB operations.
     */ protected applyPatch ( ) { this . logger . debug ( 'Patched MongoDB' ) ; if ( this . moduleExports . Server ) { this . logger . debug ( 'patching mongodb-core.Server.prototype functions: insert, remove, command, update' ) ; shimmer . wrap ( this . moduleExports . Server . prototype , 'insert' , this . getPatchCommand ( 'mongodb-insert' ) ) ; shimmer . wrap ( this . moduleExports . Server . prototype , 'remove' , this . getPatchCommand ( 'mongodb-remove' ) ) ; shimmer . wrap ( this . moduleExports . Server . prototype , 'command' , this . getPatchCommand ( 'mongodb-command' ) ) ; shimmer . wrap ( this . moduleExports . Server . prototype , 'update' , this . getPatchCommand ( 'mongodb-update' ) ) ; } if ( this . moduleExports . Cursor ) { this . logger . debug ( 'patching mongodb-core.Cursor.prototype.next' ) ; shimmer . wrap ( this . moduleExports . Cursor . prototype , 'next' , this . getPatchCursor ( ) ) ; } if ( this . internalFilesExports . ConnectionPool ) { this . logger . debug ( 'patching mongodb-core/lib/connection/pool' ) ; shimmer . wrap ( this . internalFilesExports . ConnectionPool . prototype , 'once' as never , this . getPatchEventEmitter ( ) ) ; } return this . moduleExports ; } /** Unpatches all MongoDB patched functions. */ applyUnpatch ( ) { shimmer . unwrap ( this . moduleExports . Server . prototype , 'insert' ) ; shimmer . unwrap ( this . moduleExports . Server . prototype , 'remove' ) ; shimmer . unwrap ( this . moduleExports . Server . prototype , 'command' ) ; shimmer . unwrap ( this . moduleExports . Server . prototype , 'update' ) ; shimmer . unwrap ( this . moduleExports . Cursor . prototype , 'next' ) ; if ( this . internalFilesExports . ConnectionPool ) { shimmer . unwrap ( this . internalFilesExports . ConnectionPool . prototype , 'once' ) ; } } /** Creates spans for Command operations */ private getPatchCommand ( label ) { const plugin = this ; return ( original ) => { return function ( ns , command , options , callback ) { const resultHandler = typeof options === 'function' ? options : callback ; if ( plugin . tracer . currentRootSpan && typeof resultHandler === 'function' ) { let type ; 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 . startChildSpan ( label , SpanKind . CLIENT ) ; if ( span === null ) return original . apply ( this , arguments ) ; span . addAttribute ( 'database' , ns ) ; span . addAttribute ( 'type' , type ) ; if ( plugin . options . detailedCommands === true ) { span . addAttribute ( 'command' , JSON . stringify ( command ) ) ; } if ( typeof options === 'function' ) { return original . call ( this , ns , command , plugin . patchEnd ( span , options ) ) ; } else { return original . call ( this , ns , command , options , plugin . patchEnd ( span , callback ) ) ; } } return original . apply ( this , arguments ) ; } ; } ; } /** Creates spans for Cursor operations */ private getPatchCursor ( ) { const plugin = this ; return ( original ) => { return function ( ... args ) { let resultHandler = args [ 0 ] ; if ( plugin . tracer . currentRootSpan && typeof resultHandler === 'function' ) { const span = plugin . tracer . startChildSpan ( 'mongodb-find' , SpanKind . CLIENT ) ; if ( span === null ) return original . apply ( this , arguments ) ; resultHandler = plugin . patchEnd ( span , resultHandler ) ; span . addAttribute ( 'database' , this . ns ) ; if ( plugin . options . detailedCommands === true && typeof this . cmd . query === 'object' ) { span . addAttribute ( 'command' , JSON . stringify ( this . cmd . query ) ) ; } } return original . call ( this , resultHandler ) ; } ; } ; } /** Propagate context in the event emitter of the connection pool */ private getPatchEventEmitter ( ) { const plugin = this ; return ( original ) => { return function ( event , cb ) { return original . call ( this , event , plugin . tracer . wrap ( cb ) ) ; } ; } ; } param span param resultHandler patchEnd ( span , resultHandler ) { const plugin = this ; const patchedEnd = function ( err , res ) { if ( plugin . options . detailedCommands === true && err instanceof Error ) { span . addAttribute ( 'error' , err . message ) ; } if ( span . ended === false ) { span . end ( ) ; } return resultHandler . apply ( this , arguments ) ; } ; return this . tracer . wrap ( patchedEnd ) ; } } const plugin = new MongoDBPlugin ( 'mongodb-core' ) ; export { plugin } ;