From d8c7fe47d31aae8bf2c3ce8c40eb0352d9f3e1fe Mon Sep 17 00:00:00 2001
From: Tucker Whitehouse <TuckerWhitehouse@gmail.com>
Date: Wed, 26 Aug 2020 20:39:22 -0400
Subject: [PATCH] Handle operations that return a promise instead of an event
 emitter

---
 packages/mysql/lib/mysql_p.js | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/packages/mysql/lib/mysql_p.js b/packages/mysql/lib/mysql_p.js
index 94d6c9e1..11823273 100644
--- a/packages/mysql/lib/mysql_p.js
+++ b/packages/mysql/lib/mysql_p.js
@@ -90,7 +90,7 @@ function patchOf(poolCluster) {
 
     var resultPool = poolCluster[baseFcn].apply(poolCluster, args);
     return patchObject(resultPool);
-  }
+  };
 }
 
 function patchGetConnection(pool) {
@@ -105,13 +105,13 @@ function patchGetConnection(pool) {
       args[args.length-1] = (err, connection) => {
         if(connection) patchObject(connection);
         return callback(err, connection);
-      }
+      };
     }
 
     var result = pool[baseFcn].apply(pool, args);
     if (result && result.then instanceof Function) return result.then(patchObject);
     else return result;
-  }
+  };
 }
 
 function patchObject(connection) {
@@ -208,10 +208,6 @@ function captureOperation(name) {
     command = originalOperation.call(this, sql, args.values, args.callback);
 
     if (!args.callback) {
-      command.on('end', function() {
-        subsegment.close();
-      });
-
       var errorCapturer = function (err) {
         subsegment.close(err);
 
@@ -222,14 +218,26 @@ function captureOperation(name) {
         }
       };
 
-      command.on(events.errorMonitor || 'error', errorCapturer);
+      if (command.then instanceof Function) {
+        command.then(() => {
+          subsegment.close();
+        });
+
+        command.catch(errorCapturer);
+      } else {
+        command.on('end', function() {
+          subsegment.close();
+        });
+
+        command.on(events.errorMonitor || 'error', errorCapturer);
+      }
     }
 
     subsegment.addSqlData(createSqlData(config, command));
     subsegment.namespace = 'remote';
 
     return command;
-  }
+  };
 }
 
 function createSqlData(config, command) {