From 1b0738cca56c0d819850ef481eb15a7abbe87cec Mon Sep 17 00:00:00 2001
From: Shadowbeetle <tamas.kadlecsik@gmail.com>
Date: Wed, 18 Jan 2017 14:53:47 +0100
Subject: [PATCH] fix(wrapQuery): fix when no readStream or promise is returned

---
 .gitignore                              |  4 +++
 .npmignore                              |  4 +++
 lib/instrumentations/utils/wrapQuery.js | 48 +++++++++++++++----------
 package.json                            |  2 +-
 4 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/.gitignore b/.gitignore
index f01dba8..3c633fc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,3 +29,7 @@ Desktop.ini
 .idea
 *.iml
 .vscode
+
+# RCs
+.nvmrc
+.npmrc
diff --git a/.npmignore b/.npmignore
index d73e09a..ff1bf68 100644
--- a/.npmignore
+++ b/.npmignore
@@ -44,3 +44,7 @@ Desktop.ini
 .idea
 *.iml
 .vscode
+
+#RCs
+.nvmrc
+.npmrc
diff --git a/lib/instrumentations/utils/wrapQuery.js b/lib/instrumentations/utils/wrapQuery.js
index 83919f8..fc34b9c 100644
--- a/lib/instrumentations/utils/wrapQuery.js
+++ b/lib/instrumentations/utils/wrapQuery.js
@@ -51,13 +51,13 @@ function wrapQuery (original, args, agent, params) {
   } else if (continuationMethod === 'readStream') {
     return wrapReadStream.call(this, original, args, reporter)
   } else if (continuationMethod === 'callback') { // uses callback
-    return wrapCallback.call(this, original, _params, args, reporter)
+    return wrapCallback.call(this, original, args, reporter)
   } else {
     return original.apply(this, args) // we might not want to instrument the method
   }
 }
 
-function wrapCallback (original, params, args, reporter) {
+function wrapCallback (original, args, reporter) {
   var wrappedCallback = function (original) {
     return function (err) {
       reporter.reportReceive(err)
@@ -78,33 +78,43 @@ function wrapCallback (original, params, args, reporter) {
 }
 
 function wrapPromise (original, args, reporter) {
-  reporter.reportSend()
   var originalPromise = original.apply(this, args)
-  return originalPromise.then(
-    function (v) { reporter.reportReceive(); return v },
-    function (err) { reporter.reportReceive(err); throw err }
-  )
+  if (originalPromise) {
+    reporter.reportSend()
+    return originalPromise.then(
+      function (v) {
+        reporter.reportReceive()
+        return v
+      },
+      function (err) {
+        reporter.reportReceive(err)
+        throw err
+      }
+    )
+  }
 }
 
 function wrapReadStream (original, args, reporter) {
-  reporter.reportSend()
   var originalStream = original.apply(this, args)
 
-  originalStream.on('end', function () {
-    reporter.reportReceive()
-  })
+  if (originalStream) {
+    reporter.reportSend()
+    originalStream.on('end', function () {
+      reporter.reportReceive()
+    })
 
-  originalStream.on('error', function (err) {
-    reporter.reportReceive(err)
+    originalStream.on('error', function (err) {
+      reporter.reportReceive(err)
 
-    if (typeof originalStream.listenerCount === 'function') {
-      if (originalStream.listenerCount('error') < 2) {
+      if (typeof originalStream.listenerCount === 'function') {
+        if (originalStream.listenerCount('error') < 2) {
+          throw err
+        }
+      } else if (EventEmitter.listenerCount(originalStream, 'error') < 2) {
         throw err
       }
-    } else if (EventEmitter.listenerCount(originalStream, 'error') < 2) {
-      throw err
-    }
-  })
+    })
+  }
 
   return originalStream
 }
diff --git a/package.json b/package.json
index 2c3b88d..54edcea 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@risingstack/trace",
-  "version": "2.33.0",
+  "version": "3.0.2",
   "author": "RisingStack, Inc.",
   "license": "SEE LICENSE IN LICENSE",
   "contributors": "RisingStack",