feat: added proper instrumentation for PoolNamespace.prototype.query

newrelic · May 22, 2023 · 705b01b · 705b01b
1 parent e2d8e99
commit 705b01b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 7 deletions.
diff --git a/lib/instrumentation/mysql/mysql.js b/lib/instrumentation/mysql/mysql.js
@@ -49,7 +49,8 @@ exports.callbackInitialize = function callbackInitialize(shim, mysql) {
       if (poolNamespace[symbols.clusterOf]) {
         return
       }
-      if (wrapGetConnection(shim, poolNamespace)) {
+
+      if (wrapGetConnection(shim, poolNamespace) && wrapQueryable(shim, poolNamespace, false)) {
         poolNamespace[symbols.clusterOf] = true
       }
     }

diff --git a/test/versioned/mysql/basic-pool.tap.js b/test/versioned/mysql/basic-pool.tap.js
@@ -89,7 +89,8 @@ tap.test('mysql built-in connection pools', function (t) {
   let mysql = null
   let pool = null
 
-  t.beforeEach(function () {
+  t.beforeEach(async function () {
+    await setup(require('mysql'))
     agent = helper.instrumentMockedAgent()
     mysql = require('mysql')
     pool = mysql.createPool(config)
@@ -395,7 +396,8 @@ tap.test('poolCluster', function (t) {
   let mysql = null
   let poolCluster = null
 
-  t.beforeEach(function () {
+  t.beforeEach(async function () {
+    await setup(require('mysql'))
     agent = helper.instrumentMockedAgent()
     mysql = require('mysql')
     poolCluster = mysql.createPoolCluster()
@@ -623,7 +625,7 @@ tap.test('poolCluster', function (t) {
           t.ok(transaction, 'transaction should exist')
           t.equal(transaction, txn, 'transaction must be same')
 
-          let segment = agent.tracer.getSegment().children[1]
+          let segment = agent.tracer.getSegment().parent
           t.ok(segment, 'segment should exist')
           t.ok(segment.timer.start > 0, 'starts at a positive time')
           t.ok(segment.timer.start <= Date.now(), 'starts in past')
@@ -643,7 +645,7 @@ tap.test('poolCluster', function (t) {
             t.ok(transaction, 'transaction should exist')
             t.equal(transaction, txn, 'transaction must be same')
 
-            segment = agent.tracer.getSegment().children[1]
+            segment = agent.tracer.getSegment().parent
             t.ok(segment, 'segment should exist')
             t.ok(segment.timer.start > 0, 'starts at a positive time')
             t.ok(segment.timer.start <= Date.now(), 'starts in past')

diff --git a/test/versioned/mysql/helpers.js b/test/versioned/mysql/helpers.js
@@ -7,8 +7,14 @@
 const params = require('../../lib/params')
 
 module.exports = async function setupDb(user, db, table, mysql) {
+  const regex = new RegExp(/mysql/)
   await createDb(mysql, user, db)
   await createTable(mysql, user, db, table)
+  Object.keys(require.cache).forEach((key) => {
+    if (regex.test(key)) {
+      delete require.cache[key]
+    }
+  })
 }
 
 function runCommand(client, cmd) {

diff --git a/test/versioned/mysql2/basic-pool.tap.js b/test/versioned/mysql2/basic-pool.tap.js
@@ -89,7 +89,8 @@ tap.test('mysql2 built-in connection pools', function (t) {
   let mysql = null
   let pool = null
 
-  t.beforeEach(function () {
+  t.beforeEach(async function () {
+    await setup(require('mysql2'))
     agent = helper.instrumentMockedAgent()
     mysql = require('mysql2')
     pool = mysql.createPool(config)
@@ -387,7 +388,8 @@ tap.test('poolCluster', function (t) {
   let mysql = null
   let poolCluster = null
 
-  t.beforeEach(function () {
+  t.beforeEach(async function () {
+    await setup(require('mysql2'))
     agent = helper.instrumentMockedAgent()
     mysql = require('mysql2')
     poolCluster = mysql.createPoolCluster()