@@ -51,13 +51,13 @@ function wrapQuery (original, args, agent, params) {
51
51
} else if ( continuationMethod === 'readStream' ) {
52
52
return wrapReadStream . call ( this , original , args , reporter )
53
53
} else if ( continuationMethod === 'callback' ) { // uses callback
54
- return wrapCallback . call ( this , original , _params , args , reporter )
54
+ return wrapCallback . call ( this , original , args , reporter )
55
55
} else {
56
56
return original . apply ( this , args ) // we might not want to instrument the method
57
57
}
58
58
}
59
59
60
- function wrapCallback ( original , params , args , reporter ) {
60
+ function wrapCallback ( original , args , reporter ) {
61
61
var wrappedCallback = function ( original ) {
62
62
return function ( err ) {
63
63
reporter . reportReceive ( err )
@@ -78,33 +78,43 @@ function wrapCallback (original, params, args, reporter) {
78
78
}
79
79
80
80
function wrapPromise ( original , args , reporter ) {
81
- reporter . reportSend ( )
82
81
var originalPromise = original . apply ( this , args )
83
- return originalPromise . then (
84
- function ( v ) { reporter . reportReceive ( ) ; return v } ,
85
- function ( err ) { reporter . reportReceive ( err ) ; throw err }
86
- )
82
+ if ( originalPromise ) {
83
+ reporter . reportSend ( )
84
+ return originalPromise . then (
85
+ function ( v ) {
86
+ reporter . reportReceive ( )
87
+ return v
88
+ } ,
89
+ function ( err ) {
90
+ reporter . reportReceive ( err )
91
+ throw err
92
+ }
93
+ )
94
+ }
87
95
}
88
96
89
97
function wrapReadStream ( original , args , reporter ) {
90
- reporter . reportSend ( )
91
98
var originalStream = original . apply ( this , args )
92
99
93
- originalStream . on ( 'end' , function ( ) {
94
- reporter . reportReceive ( )
95
- } )
100
+ if ( originalStream ) {
101
+ reporter . reportSend ( )
102
+ originalStream . on ( 'end' , function ( ) {
103
+ reporter . reportReceive ( )
104
+ } )
96
105
97
- originalStream . on ( 'error' , function ( err ) {
98
- reporter . reportReceive ( err )
106
+ originalStream . on ( 'error' , function ( err ) {
107
+ reporter . reportReceive ( err )
99
108
100
- if ( typeof originalStream . listenerCount === 'function' ) {
101
- if ( originalStream . listenerCount ( 'error' ) < 2 ) {
109
+ if ( typeof originalStream . listenerCount === 'function' ) {
110
+ if ( originalStream . listenerCount ( 'error' ) < 2 ) {
111
+ throw err
112
+ }
113
+ } else if ( EventEmitter . listenerCount ( originalStream , 'error' ) < 2 ) {
102
114
throw err
103
115
}
104
- } else if ( EventEmitter . listenerCount ( originalStream , 'error' ) < 2 ) {
105
- throw err
106
- }
107
- } )
116
+ } )
117
+ }
108
118
109
119
return originalStream
110
120
}
0 commit comments