fix(query): preserve cls context for logger (#12328)

lib/dialects/mssql/query.js

@@ -44,48 +44,29 @@ class Query extends AbstractQuery {
     return paramType;
   }
 
-  _run(connection, sql, parameters) {
+  async _run(connection, sql, parameters) {
     this.sql = sql;
     const { options } = this;
 
     const complete = this._logQuery(sql, debug, parameters);
 
-    return new Promise((resolve, reject) => {
-      const handleTransaction = err => {
-        if (err) {
-          err.sql = sql;
-          err.parameters = parameters;
-          reject(this.formatError(err));
-          return;
-        }
-        resolve(this.formatResults());
-      };
+    const query = new Promise((resolve, reject) => {
       // TRANSACTION SUPPORT
       if (sql.startsWith('BEGIN TRANSACTION')) {
-        return connection.beginTransaction(handleTransaction, options.transaction.name, connection.lib.ISOLATION_LEVEL[options.isolationLevel]);
+        return connection.beginTransaction(error => error ? reject(error) : resolve([]), options.transaction.name, connection.lib.ISOLATION_LEVEL[options.isolationLevel]);
       }
       if (sql.startsWith('COMMIT TRANSACTION')) {
-        return connection.commitTransaction(handleTransaction);
+        return connection.commitTransaction(error => error ? reject(error) : resolve([]));
       }
       if (sql.startsWith('ROLLBACK TRANSACTION')) {
-        return connection.rollbackTransaction(handleTransaction, options.transaction.name);
+        return connection.rollbackTransaction(error => error ? reject(error) : resolve([]), options.transaction.name);
       }
       if (sql.startsWith('SAVE TRANSACTION')) {
-        return connection.saveTransaction(handleTransaction, options.transaction.name);
+        return connection.saveTransaction(error => error ? reject(error) : resolve([]), options.transaction.name);
       }
-      const results = [];
-      const request = new connection.lib.Request(sql, (err, rowCount) => {
 
-        complete();
-
-        if (err) {
-          err.sql = sql;
-          err.parameters = parameters;
-          reject(this.formatError(err));
-        } else {
-          resolve(this.formatResults(results, rowCount));
-        }
-      });
+      const rows = [];
+      const request = new connection.lib.Request(sql, (err, rowCount) => err ? reject(err) : resolve([rows, rowCount]));
 
       if (parameters) {
         _.forOwn(parameters, (value, key) => {
@@ -95,6 +76,27 @@ class Query extends AbstractQuery {
       }
 
       request.on('row', columns => {
+        rows.push(columns);
+      });
+
+      connection.execSql(request);
+    });
+
+    let rows, rowCount;
+
+    try {
+      [rows, rowCount] = await query;
+    } catch (err) {
+      err.sql = sql;
+      err.parameters = parameters;
+
+      throw this.formatError(err);
+    }
+
+    complete();
+
+    if (Array.isArray(rows)) {
+      rows = rows.map(columns => {
         const row = {};
         for (const column of columns) {
           const typeid = column.metadata.type.id;
@@ -106,12 +108,11 @@ class Query extends AbstractQuery {
           }
           row[column.metadata.colName] = value;
         }
-
-        results.push(row);
+        return row;
       });
+    }
 
-      connection.execSql(request);
-    });
+    return this.formatResults(rows, rowCount);
   }
 
   run(sql, parameters) {

lib/dialects/mysql/query.js

@@ -35,11 +35,15 @@ class Query extends AbstractQuery {
 
     const complete = this._logQuery(sql, debug, parameters);
 
-    const results = await new Promise((resolve, reject) => {
-      const handler = (err, results) => {
-        complete();
+    const query = parameters && parameters.length
+      ? new Promise((resolve, reject) => connection.execute(sql, parameters, (error, result) => error ? reject(error) : resolve(result)).setMaxListeners(100))
+      : new Promise((resolve, reject) => connection.query({ sql }, (error, result) => error ? reject(error) : resolve(result)).setMaxListeners(100));
+
+    let results;
 
-        if (err) {
+    try {
+      results = await query;
+    } catch (err) {
       // MySQL automatically rolls-back transactions in the event of a deadlock
       if (options.transaction && err.errno === 1213) {
         options.transaction.finished = 'rollback';
@@ -47,18 +51,11 @@ class Query extends AbstractQuery {
       err.sql = sql;
       err.parameters = parameters;
 
-          reject(this.formatError(err));
-        } else {
-          resolve(results);
-        }
-      };
-      if (parameters) {
-        debug('parameters(%j)', parameters);
-        connection.execute(sql, parameters, handler).setMaxListeners(100);
-      } else {
-        connection.query({ sql }, handler).setMaxListeners(100);
+      throw this.formatError(err);
     }
-    });
+
+    complete();
+
     // Log warnings if we've got them.
     if (showWarnings && results && results.warningStatus > 0) {
       await this.logWarnings(results);

test/integration/cls.test.js

@@ -6,7 +6,8 @@ const chai      = require('chai'),
   Sequelize = Support.Sequelize,
   cls       = require('cls-hooked'),
   current = Support.sequelize,
-  delay     = require('delay');
+  delay     = require('delay'),
+  sinon = require('sinon');
 
 if (current.dialect.supports.transactions) {
   describe(Support.getTestDialectTeaser('CLS (Async hooks)'), () => {
@@ -146,5 +147,35 @@ if (current.dialect.supports.transactions) {
       }
       );
     });
+
+    it('custom logging with benchmarking has correct CLS context', async function() {
+      const logger = sinon.spy(() => {
+        return this.ns.get('value');
+      });
+      const sequelize = Support.createSequelizeInstance({
+        logging: logger,
+        benchmark: true
+      });
+
+      const result = this.ns.runPromise(async () => {
+        this.ns.set('value', 1);
+        await delay(500);
+        return sequelize.query('select 1;');
+      });
+
+      await this.ns.runPromise(() => {
+        this.ns.set('value', 2);
+        return sequelize.query('select 2;');
+      });
+
+      await result;
+
+      expect(logger.calledTwice).to.be.true;
+      expect(logger.firstCall.args[0]).to.be.match(/Executed \((\d*|default)\): select 2/);
+      expect(logger.firstCall.returnValue).to.be.equal(2);
+      expect(logger.secondCall.args[0]).to.be.match(/Executed \((\d*|default)\): select 1/);
+      expect(logger.secondCall.returnValue).to.be.equal(1);
+
+    });
   });
 }