Add support for including sql query in sql subsegment for MySQL (#564)

* Add support for including sql query in sql subsegment for MySQL * Update createSqlData to accept values and sql as arguments, and the call to createSqlData to send those values from the arguments made to the sql call * Add function comments to mysqL_p.createSqlData --------- Co-authored-by: Jason Terando <jason.t@digitalroominc.com>
aws · Mar 7, 2023 · 3abe727 · 3abe727
1 parent 0fe39de
commit 3abe727
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 13 deletions.
diff --git a/packages/core/README.md b/packages/core/README.md
@@ -59,7 +59,7 @@ section for different usages.
     AWS_XRAY_DAEMON_ADDRESS          For setting the daemon address and port.
     AWS_XRAY_CONTEXT_MISSING         For setting the SDK behavior when trace context is missing. Valid values are 'RUNTIME_ERROR', 'IGNORE_ERROR' or 'LOG_ERROR'. The SDK's default behavior is 'LOG_ERROR'.
     AWS_XRAY_LOG_LEVEL               Sets a log level for the SDK built in logger. This value is ignored if AWS_XRAY_DEBUG_MODE is set.
-    AWS_XRAY_COLLECT_SQL_QUERIES     Enables SQL query capture (currently only Postgres supported)
+    AWS_XRAY_COLLECT_SQL_QUERIES     Enables SQL query capture (currently only Postgres and MySQL supported)
 
 ### Daemon configuration
 

diff --git a/packages/mysql/lib/mysql_p.js b/packages/mysql/lib/mysql_p.js
@@ -242,9 +242,7 @@ function captureOperation(name) {
       if (isPromise(command)) {
         command.then(() => {
           subsegment.close();
-        });
-
-        command.catch(errorCapturer);
+        }).catch (errorCapturer);
       } else {
         command.on('end', function() {
           subsegment.close();
@@ -254,19 +252,32 @@ function captureOperation(name) {
       }
     }
 
-    subsegment.addSqlData(createSqlData(config, command));
+    subsegment.addSqlData(createSqlData(config, args.values, args.sql));
     subsegment.namespace = 'remote';
 
     return command;
   };
 }
 
-function createSqlData(config, command) {
-  var commandType = command.values ? PREPARED : null;
-
+/**
+ * Generate a SQL data object.  Note that this implementation differs from
+ * that in postgres_p.js because the posgres client structures commands
+ * and prepared statements differently than mysql/mysql2.
+ *
+ * @param {object} config
+ * @param {any} values
+ * @param {string} sql
+ * @returns SQL data object
+ */
+function createSqlData(config, values, sql) {
+  var commandType = values ? PREPARED : null;
   var data = new SqlData(DATABASE_VERS, DRIVER_VERS, config.user,
     config.host + ':' + config.port + '/' + config.database,
     commandType);
 
+  if (process.env.AWS_XRAY_COLLECT_SQL_QUERIES && sql) {
+    data.sanitized_query = sql;
+  }
+
   return data;
 }
diff --git a/packages/mysql/test/unit/mysql_p.test.js b/packages/mysql/test/unit/mysql_p.test.js
@@ -141,7 +141,7 @@ describe('captureMySQL', function() {
         var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
         var config = conn.config;
 
-        query.call(connectionObj, 'sql here');
+        query.call(connectionObj, 'sql here', [1]);
 
         stubDataInit.should.have.been.calledWithExactly(undefined, undefined, config.user,
           config.host + ':' + config.port + '/' + config.database, 'statement');
@@ -247,6 +247,34 @@ describe('captureMySQL', function() {
           done();
         }, 50);
       });
+      it('should add query to the subsegments sql data when AWS_XRAY_COLLECT_SQL_QUERIES is truthy', function () {
+        sandbox.stub(process, 'env').value({ ...AWSXRay, 'AWS_XRAY_COLLECT_SQL_QUERIES': true });
+        var stubAddSql = sandbox.stub(subsegment, 'addSqlData');
+        var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
+        var conParam = connectionObj.config;
+
+        query.call(connectionObj, 'sql here', [1]);
+
+        stubDataInit.should.have.been.calledWithExactly(process.env.MYSQL_DATABASE_VERSION, process.env.MYSQL_DRIVER_VERSION,
+          conParam.user, conParam.host + ':' + conParam.port + '/' + conParam.database, 'statement');
+        stubAddSql.should.have.been.calledWithExactly(sinon.match.instanceOf(SqlData));
+        stubAddSql.should.have.been.calledWithExactly(sinon.match.has('sanitized_query', 'sql here'));
+      });
+      it('should NOT add query to the subsegments sql data when AWS_XRAY_COLLECT_SQL_QUERIES is not set', function () {
+        sandbox.stub(process, 'env').value({ ...AWSXRay, 'AWS_XRAY_COLLECT_SQL_QUERIES': undefined });
+        var stubAddSql = sandbox.stub(subsegment, 'addSqlData');
+        var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
+        var conParam = connectionObj.config;
+
+        query.call(connectionObj, 'sql here', [1]);
+
+        stubDataInit.should.have.been.calledWithExactly(process.env.MYSQL_DATABASE_VERSION, process.env.MYSQL_DRIVER_VERSION,
+          conParam.user, conParam.host + ':' + conParam.port + '/' + conParam.database, 'statement');
+        stubAddSql.should.have.been.calledWithExactly(sinon.match.instanceOf(SqlData));
+        sinon.assert.match(sinon.match, {
+          'sanitized_query': undefined
+        });
+      });
     });
   });
 
@@ -344,7 +372,6 @@ describe('captureMySQL', function() {
           stubClose.should.have.been.calledWithExactly(err);
         });
       });
-
     });
   });
 
@@ -430,7 +457,7 @@ describe('captureMySQL', function() {
         var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
         var config = conn.config;
 
-        query.call(connectionObj, 'sql here');
+        query.call(connectionObj, 'sql here', [1]);
 
         stubDataInit.should.have.been.calledWithExactly(undefined, undefined, config.user,
           config.host + ':' + config.port + '/' + config.database, 'statement');
@@ -526,7 +553,7 @@ describe('captureMySQL', function() {
         var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
         var config = conn.config;
 
-        query.call(connectionObj, 'sql here');
+        query.call(connectionObj, 'sql here', [1]);
 
         stubDataInit.should.have.been.calledWithExactly(undefined, undefined, config.user,
           config.host + ':' + config.port + '/' + config.database, 'statement');
@@ -673,7 +700,7 @@ describe('captureMySQL', function() {
         var stubDataInit = sandbox.stub(SqlData.prototype, 'init');
         var config = conn.config;
 
-        query.call(connectionObj, 'sql here');
+        query.call(connectionObj, 'sql here', [1]);
 
         stubDataInit.should.have.been.calledWithExactly(undefined, undefined, config.user,
           config.host + ':' + config.port + '/' + config.database, 'statement');