import { channel , IModulePatcher , PatchFunction } from "diagnostic-channel" ; import { EventEmitter } from "events" ; export interface IPostgresResult { rowCount : number ; command : string ; } export interface IPostgresData { query : { text ? : string ; plan ? : string ; preparable ? : { text : string ; args : any [ ] ; } ; } ; database : { host : string ; port : string ; } ; result ? : IPostgresResult ; duration : number ; error ? : Error ; time : Date ; } type PostgresCallback = ( err : Error , res : IPostgresResult ) => any ; function postgres6PatchFunction ( originalPg , originalPgPath ) { const originalClientQuery = originalPg . Client . prototype . query ; const diagnosticOriginalFunc = "__diagnosticOriginalFunc" ; originalPg . Client . prototype . query = function query ( config , values , callback ) { const data = { query : { } , database : { host : this . connectionParameters . host , port : this . connectionParameters . port , } , result : null , error : null , duration : 0 , time : new Date ( ) , } ; const start = process . hrtime ( ) ; let queryResult ; function patchCallback ( cb ? ) { if ( cb && cb [ diagnosticOriginalFunc ] ) { cb = cb [ diagnosticOriginalFunc ] ; } const trackingCallback = channel . bindToContext ( function ( err , res ) : any { const end = process . hrtime ( start ) ; data . result = res && { rowCount : res . rowCount , command : res . command } ; data . error = err ; data . duration = Math . ceil ( ( end [ 0 ] * 1e3 ) + ( end [ 1 ] / 1e6 ) ) ; channel . publish ( "postgres" , data ) ; if ( err ) { if ( cb ) { return cb . apply ( this , arguments ) ; } else if ( queryResult && queryResult instanceof EventEmitter ) { queryResult . emit ( "error" , err ) ; } } else if ( cb ) { cb . apply ( this , arguments ) ; } } ) ; try { Object . defineProperty ( trackingCallback , diagnosticOriginalFunc , { value : cb } ) ; return trackingCallback ; } catch ( e ) { return cb ; } } try { if ( typeof config === "string" ) { if ( values instanceof Array ) { data . query . preparable = { text : config , args : values , } ; callback = patchCallback ( callback ) ; } else { data . query . text = config ; if ( callback ) { callback = patchCallback ( callback ) ; } else { values = patchCallback ( values ) ; } } } else { if ( typeof config . name === "string" ) { data . query . plan = config . name ; } else if ( config . values instanceof Array ) { data . query . preparable = { text : config . text , args : config . values , } ; } else { data . query . text = config . text ; } if ( callback ) { callback = patchCallback ( callback ) ; } else if ( values ) { values = patchCallback ( values ) ; } else { config . callback = patchCallback ( config . callback ) ; } } } catch ( e ) { return originalClientQuery . apply ( this , arguments ) ; } arguments [ 0 ] = config ; arguments [ 1 ] = values ; arguments [ 2 ] = callback ; arguments . length = ( arguments . length > 3 ) ? arguments . length : 3 ; queryResult = originalClientQuery . apply ( this , arguments ) ; return queryResult ; } ; return originalPg ; } function postgres7PatchFunction ( originalPg , originalPgPath ) { const originalClientQuery = originalPg . Client . prototype . query ; const diagnosticOriginalFunc = "__diagnosticOriginalFunc" ; originalPg . Client . prototype . query = function query ( config , values , callback ) { let callbackProvided = ! ! callback ; const data = { query : { } , database : { host : this . connectionParameters . host , port : this . connectionParameters . port , } , result : null , error : null , duration : 0 , time : new Date ( ) , } ; const start = process . hrtime ( ) ; let queryResult ; function patchCallback ( cb ? ) { if ( cb && cb [ diagnosticOriginalFunc ] ) { cb = cb [ diagnosticOriginalFunc ] ; } const trackingCallback = channel . bindToContext ( function ( err , res ) : any { const end = process . hrtime ( start ) ; data . result = res && { rowCount : res . rowCount , command : res . command } ; data . error = err ; data . duration = Math . ceil ( ( end [ 0 ] * 1e3 ) + ( end [ 1 ] / 1e6 ) ) ; channel . publish ( "postgres" , data ) ; if ( err ) { if ( cb ) { return cb . apply ( this , arguments ) ; } else if ( queryResult && queryResult instanceof EventEmitter ) { queryResult . emit ( "error" , err ) ; } } else if ( cb ) { cb . apply ( this , arguments ) ; } } ) ; try { Object . defineProperty ( trackingCallback , diagnosticOriginalFunc , { value : cb } ) ; return trackingCallback ; } catch ( e ) { return cb ; } } try { if ( typeof config === "string" ) { if ( values instanceof Array ) { data . query . preparable = { text : config , args : values , } ; callbackProvided = typeof callback === "function" ; callback = callbackProvided ? patchCallback ( callback ) : callback ; } else { data . query . text = config ; if ( callback ) { callbackProvided = typeof callback === "function" ; callback = callbackProvided ? patchCallback ( callback ) : callback ; } else { callbackProvided = typeof values === "function" ; values = callbackProvided ? patchCallback ( values ) : values ; } } } else { if ( typeof config . name === "string" ) { data . query . plan = config . name ; } else if ( config . values instanceof Array ) { data . query . preparable = { text : config . text , args : config . values , } ; } else { data . query . text = config . text ; } if ( callback ) { callbackProvided = typeof callback === "function" ; callback = patchCallback ( callback ) ; } else if ( values ) { callbackProvided = typeof values === "function" ; values = callbackProvided ? patchCallback ( values ) : values ; } else { callbackProvided = typeof config . callback === "function" ; config . callback = callbackProvided ? patchCallback ( config . callback ) : config . callback ; } } } catch ( e ) { return originalClientQuery . apply ( this , arguments ) ; } arguments [ 0 ] = config ; arguments [ 1 ] = values ; arguments [ 2 ] = callback ; arguments . length = ( arguments . length > 3 ) ? arguments . length : 3 ; queryResult = originalClientQuery . apply ( this , arguments ) ; if ( ! callbackProvided ) { return queryResult . then ( ( result ) => { patchCallback ( ) ( undefined , result ) ; return new this . _Promise ( ( resolve , reject ) => { resolve ( result ) ; } ) ; } ) . catch ( ( error ) => { patchCallback ( ) ( error , undefined ) ; return new this . _Promise ( ( resolve , reject ) => { reject ( error ) ; } ) ; } ) ; } return queryResult ; } ; return originalPg ; } export const postgres6 = { versionSpecifier : "6.*" , patch : postgres6PatchFunction , } ; export const postgres7 = { versionSpecifier : ">=7.* <=8.*" , patch : postgres7PatchFunction , } ; export function enable ( ) { channel . registerMonkeyPatch ( "pg" , postgres6 ) ; channel . registerMonkeyPatch ( "pg" , postgres7 ) ; }