From 268adf413185374135225a54d3e7d96a2af244a4 Mon Sep 17 00:00:00 2001
From: Michele Azzolari <michele@azzolari.it>
Date: Thu, 16 Jan 2025 17:27:27 +0100
Subject: [PATCH] fix(instrumentation-mysql2): patch Connection class when
 imported (#2590)

---
 .../src/instrumentation.ts                    | 88 +++++++++++++------
 .../src/utils.ts                              | 13 +--
 2 files changed, 67 insertions(+), 34 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts
index 8c90080980..4996cdec13 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql2/src/instrumentation.ts
@@ -18,6 +18,7 @@ import * as api from '@opentelemetry/api';
 import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
+  InstrumentationNodeModuleFile,
   isWrapped,
   safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
@@ -41,6 +42,8 @@ import { PACKAGE_NAME, PACKAGE_VERSION } from './version';
 
 type formatType = typeof mysqlTypes.format;
 
+const supportedVersions = ['>=1.4.2 <4'];
+
 export class MySQL2Instrumentation extends InstrumentationBase<MySQL2InstrumentationConfig> {
   static readonly COMMON_ATTRIBUTES = {
     [SEMATTRS_DB_SYSTEM]: DBSYSTEMVALUES_MYSQL,
@@ -51,45 +54,74 @@ export class MySQL2Instrumentation extends InstrumentationBase<MySQL2Instrumenta
   }
 
   protected init() {
+    let format: formatType | undefined;
+    const patch = (ConnectionPrototype: mysqlTypes.Connection) => {
+      if (isWrapped(ConnectionPrototype.query)) {
+        this._unwrap(ConnectionPrototype, 'query');
+      }
+      this._wrap(
+        ConnectionPrototype,
+        'query',
+        this._patchQuery(format, false) as any
+      );
+      if (isWrapped(ConnectionPrototype.execute)) {
+        this._unwrap(ConnectionPrototype, 'execute');
+      }
+      this._wrap(
+        ConnectionPrototype,
+        'execute',
+        this._patchQuery(format, true) as any
+      );
+    };
+    const unpatch = (ConnectionPrototype: mysqlTypes.Connection) => {
+      this._unwrap(ConnectionPrototype, 'query');
+      this._unwrap(ConnectionPrototype, 'execute');
+    };
     return [
       new InstrumentationNodeModuleDefinition(
         'mysql2',
-        ['>=1.4.2 <4'],
+        supportedVersions,
         (moduleExports: any) => {
-          const ConnectionPrototype: mysqlTypes.Connection =
-            getConnectionPrototypeToInstrument(moduleExports.Connection);
-          if (isWrapped(ConnectionPrototype.query)) {
-            this._unwrap(ConnectionPrototype, 'query');
-          }
-          this._wrap(
-            ConnectionPrototype,
-            'query',
-            this._patchQuery(moduleExports.format, false) as any
-          );
-
-          if (isWrapped(ConnectionPrototype.execute)) {
-            this._unwrap(ConnectionPrototype, 'execute');
+          if (!format && moduleExports.format) {
+            format = moduleExports.format;
           }
-          this._wrap(
-            ConnectionPrototype,
-            'execute',
-            this._patchQuery(moduleExports.format, true) as any
-          );
-
           return moduleExports;
         },
-        (moduleExports: any) => {
-          if (moduleExports === undefined) return;
-          const ConnectionPrototype: mysqlTypes.Connection =
-            moduleExports.Connection.prototype;
-          this._unwrap(ConnectionPrototype, 'query');
-          this._unwrap(ConnectionPrototype, 'execute');
-        }
+        () => {},
+        [
+          new InstrumentationNodeModuleFile(
+            'mysql2/promise.js',
+            supportedVersions,
+            (moduleExports: any) => {
+              if (!format && moduleExports.format) {
+                format = moduleExports.format;
+              }
+              return moduleExports;
+            },
+            () => {}
+          ),
+          new InstrumentationNodeModuleFile(
+            'mysql2/lib/connection.js',
+            supportedVersions,
+            (moduleExports: any) => {
+              const ConnectionPrototype: mysqlTypes.Connection =
+                getConnectionPrototypeToInstrument(moduleExports);
+              patch(ConnectionPrototype);
+              return moduleExports;
+            },
+            (moduleExports: any) => {
+              if (moduleExports === undefined) return;
+              const ConnectionPrototype: mysqlTypes.Connection =
+                getConnectionPrototypeToInstrument(moduleExports);
+              unpatch(ConnectionPrototype);
+            }
+          ),
+        ]
       ),
     ];
   }
 
-  private _patchQuery(format: formatType, isPrepared: boolean) {
+  private _patchQuery(format: formatType | undefined, isPrepared: boolean) {
     return (originalQuery: Function): Function => {
       const thisPlugin = this;
       return function query(
diff --git a/plugins/node/opentelemetry-instrumentation-mysql2/src/utils.ts b/plugins/node/opentelemetry-instrumentation-mysql2/src/utils.ts
index bb0fb763d3..9e495a7d2c 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql2/src/utils.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql2/src/utils.ts
@@ -22,6 +22,9 @@ import {
   SEMATTRS_NET_PEER_NAME,
   SEMATTRS_NET_PEER_PORT,
 } from '@opentelemetry/semantic-conventions';
+import type * as mysqlTypes from 'mysql2';
+
+type formatType = typeof mysqlTypes.format;
 
 /*
   Following types declare an expectation on mysql2 types and define a subset we
@@ -103,14 +106,12 @@ function getJDBCString(
  */
 export function getDbStatement(
   query: string | Query | QueryOptions,
-  format: (
-    sql: string,
-    values: any[],
-    stringifyObjects?: boolean,
-    timeZone?: string
-  ) => string,
+  format?: formatType,
   values?: any[]
 ): string {
+  if (!format) {
+    return typeof query === 'string' ? query : query.sql;
+  }
   if (typeof query === 'string') {
     return values ? format(query, values) : query;
   } else {