ported mysql2/promise instrumentation + tests from the external @newrelic/mysql repo

lib/instrumentation/mysql.js → lib/instrumentation/mysql/mysql.js

@@ -5,17 +5,17 @@
 
 'use strict'
 
-const dbutils = require('../db/utils')
-const properties = require('../util/properties')
+const dbutils = require('../../db/utils')
+const properties = require('../../util/properties')
 
-module.exports = function initialize(agent, mysql, moduleName, shim) {
+exports.callbackInitialize = function callbackInitialize(shim, mysql) {
   shim.setDatastore(shim.MYSQL)
   shim.__wrappedPoolConnection = false
 
   shim.wrapReturn(mysql, 'createConnection', wrapCreateConnection)
   function wrapCreateConnection(shim, fn, fnName, connection) {
     shim.logger.debug('Wrapping Connection#query')
-    if (wrapQueriable(shim, connection, false)) {
+    if (wrapQueryable(shim, connection, false)) {
       const connProto = Object.getPrototypeOf(connection)
       shim.setInternalProperty(connProto, '__NR_storeDatabase', true)
       shim.unwrap(mysql, 'createConnection')
@@ -25,7 +25,7 @@ module.exports = function initialize(agent, mysql, moduleName, shim) {
   shim.wrapReturn(mysql, 'createPool', wrapCreatePool)
   function wrapCreatePool(shim, fn, fnName, pool) {
     shim.logger.debug('Wrapping Pool#query and Pool#getConnection')
-    if (wrapQueriable(shim, pool, true) && wrapGetConnection(shim, pool)) {
+    if (wrapQueryable(shim, pool, true) && wrapGetConnection(shim, pool)) {
       shim.unwrap(mysql, 'createPool')
     }
   }
@@ -48,6 +48,13 @@ module.exports = function initialize(agent, mysql, moduleName, shim) {
   }
 }
 
+exports.promiseInitialize = function promiseInitialize(shim) {
+  const callbackAPI = shim.require('./index')
+  if (callbackAPI && !shim.isWrapped(callbackAPI.createConnection)) {
+    exports.callbackInitialize(shim, callbackAPI)
+  }
+}
+
 function wrapGetConnection(shim, connectable) {
   if (!connectable || !connectable.getConnection || shim.isWrapped(connectable.getConnection)) {
     shim.logger.trace(
@@ -91,7 +98,7 @@ function wrapGetConnectionCallback(shim, cb) {
   return function wrappedGetConnectionCallback(err, conn) {
     try {
       shim.logger.debug('Wrapping PoolConnection#query')
-      if (!err && wrapQueriable(shim, conn, false)) {
+      if (!err && wrapQueryable(shim, conn, false)) {
         // Leave getConnection wrapped in order to maintain TX state, but we can
         // simplify the wrapping of its callback in future calls.
         shim.__wrappedPoolConnection = true
@@ -106,20 +113,20 @@ function wrapGetConnectionCallback(shim, cb) {
   }
 }
 
-function wrapQueriable(shim, queriable, isPoolQuery) {
-  if (!queriable || !queriable.query || shim.isWrapped(queriable.query)) {
+function wrapQueryable(shim, queryable, isPoolQuery) {
+  if (!queryable || !queryable.query || shim.isWrapped(queryable.query)) {
     shim.logger.debug(
       {
-        queriable: !!queriable,
-        query: !!(queriable && queriable.query),
-        isWrapped: !!(queriable && shim.isWrapped(queriable.query))
+        queryable: !!queryable,
+        query: !!(queryable && queryable.query),
+        isWrapped: !!(queryable && shim.isWrapped(queryable.query))
       },
-      'Not wrappying queriable'
+      'Not wrapping queryable'
     )
     return false
   }
 
-  const proto = Object.getPrototypeOf(queriable)
+  const proto = Object.getPrototypeOf(queryable)
 
   let describe
   if (isPoolQuery) {
@@ -131,7 +138,7 @@ function wrapQueriable(shim, queriable, isPoolQuery) {
 
   shim.recordQuery(proto, 'query', describe)
 
-  if (queriable.execute) {
+  if (queryable.execute) {
     shim.recordQuery(proto, 'execute', describe)
   }
 

lib/instrumentation/mysql/nr-hooks.js

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2022 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const instrumentation = require('./mysql')
+
+/**
+ * We only need to register the instrumentation once for both mysql and mysql2
+ *  because there is some 🪄 in shimmer
+ * See: https://github.com/newrelic/node-newrelic/blob/main/lib/shimmer.js#L459
+ */
+module.exports = [
+  {
+    type: 'datastore',
+    moduleName: 'mysql',
+    onRequire: instrumentation.callbackInitialize
+  },
+  {
+    type: 'datastore',
+    moduleName: 'mysql2/promise',
+    onRequire: instrumentation.promiseInitialize
+  }
+]

lib/instrumentations.js

@@ -25,7 +25,7 @@ module.exports = function instrumentations() {
     'koa': { module: '@newrelic/koa' },
     'memcached': { type: MODULE_TYPE.DATASTORE },
     'mongodb': { type: MODULE_TYPE.DATASTORE },
-    'mysql': { type: MODULE_TYPE.DATASTORE },
+    'mysql': { module: './instrumentation/mysql' },
     'pg': { type: MODULE_TYPE.DATASTORE },
     'q': { type: null },
     'redis': { type: MODULE_TYPE.DATASTORE },

test/versioned/mysql2/promises.tap.js

@@ -0,0 +1,130 @@
+/*
+ * Copyright 2020 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const setup = require('./setup')
+const tap = require('tap')
+const helper = require('../../lib/agent_helper')
+const params = require('../../lib/params')
+const urltils = require('../../../lib/util/urltils')
+
+const { USER, DATABASE } = setup
+
+tap.test('mysql2 promises', { timeout: 30000 }, (t) => {
+  t.autoend()
+
+  let mysql = null
+  let client = null
+  let agent = null
+
+  t.beforeEach(async () => {
+    await setup(require('mysql2'))
+    agent = helper.instrumentMockedAgent()
+
+    mysql = require('mysql2/promise')
+
+    client = await mysql.createConnection({
+      user: USER,
+      database: DATABASE,
+      host: params.mysql_host,
+      port: params.mysql_port
+    })
+  })
+
+  t.afterEach(async () => {
+    helper.unloadAgent(agent)
+    if (client) {
+      await client.end()
+      client = null
+    }
+  })
+
+  t.test('basic transaction', (t) => {
+    return helper
+      .runInTransaction(agent, () => {
+        return client.query('SELECT 1').then(() => {
+          t.ok(agent.getTransaction(), 'we should be in a transaction')
+          agent.getTransaction().end()
+        })
+      })
+      .then(() => checkQueries(t, agent))
+  })
+
+  t.test('query with values', (t) => {
+    return helper
+      .runInTransaction(agent, () => {
+        return client.query('SELECT 1', []).then(() => {
+          t.ok(agent.getTransaction(), 'we should be in a transaction')
+          agent.getTransaction().end()
+        })
+      })
+      .then(() => checkQueries(t, agent))
+  })
+
+  t.test('database name should change with use statement', (t) => {
+    return helper
+      .runInTransaction(agent, () => {
+        return client
+          .query('create database if not exists test_db')
+          .then(() => {
+            t.ok(agent.getTransaction(), 'we should be in a transaction')
+            return client.query('use test_db')
+          })
+          .then(() => {
+            t.ok(agent.getTransaction(), 'we should be in a transaction')
+            return client.query('SELECT 1 + 1 AS solution')
+          })
+          .then(() => {
+            t.ok(agent.getTransaction(), 'we should be in a transaction')
+
+            const segment = agent.getTransaction().trace.root.children[2]
+            const attributes = segment.getAttributes()
+            t.equal(
+              attributes.host,
+              urltils.isLocalhost(params.mysql_host)
+                ? agent.config.getHostnameSafe()
+                : params.mysql_host,
+              'should set host name'
+            )
+            t.equal(attributes.database_name, 'test_db', 'should follow use statement')
+            t.equal(attributes.port_path_or_id, '3306', 'should set port')
+
+            agent.getTransaction().end()
+          })
+      })
+      .then(() => checkQueries(t, agent))
+  })
+
+  t.test('query with options object rather than sql', (t) => {
+    return helper
+      .runInTransaction(agent, () => {
+        return client.query({ sql: 'SELECT 1' })
+      })
+      .then(() => {
+        agent.getTransaction().end()
+      })
+      .then(() => checkQueries(t, agent))
+  })
+
+  t.test('query with options object and values', (t) => {
+    return helper
+      .runInTransaction(agent, () => {
+        return client.query({ sql: 'SELECT 1' }, [])
+      })
+      .then(() => {
+        agent.getTransaction().end()
+      })
+      .then(() => checkQueries(t, agent))
+  })
+})
+
+function checkQueries(t, agent) {
+  const querySamples = agent.queries.samples
+  t.ok(querySamples.size > 0, 'there should be a query sample')
+  for (const sample of querySamples.values()) {
+    t.ok(sample.total > 0, 'the samples should have positive duration')
+  }
+}