fix(mysql): address review comments from @Flarna, @blumamir

Signed-off-by: Simon Stone <Simon.Stone@docusign.com>
open-telemetry · Aug 11, 2021 · 5f192b8 · 5f192b8
1 parent 25b93b1
commit 5f192b8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 26 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
@@ -34,6 +34,8 @@ import { VERSION } from './version';
 
 type formatType = typeof mysqlTypes.format;
 
+type getConnectionCallbackType = (err: mysqlTypes.MysqlError, connection: mysqlTypes.PoolConnection) => void;
+
 export class MySQLInstrumentation extends InstrumentationBase<
   typeof mysqlTypes
 > {
@@ -182,21 +184,21 @@ export class MySQLInstrumentation extends InstrumentationBase<
 
         if (arguments.length === 1 && typeof arg1 === 'function') {
           const patchFn = thisPlugin._getConnectionCallbackPatchFn(
-            arg1,
+            arg1 as getConnectionCallbackType,
             format
           );
           return originalGetConnection.call(pool, patchFn);
         }
         if (arguments.length === 2 && typeof arg2 === 'function') {
           const patchFn = thisPlugin._getConnectionCallbackPatchFn(
-            arg2,
+            arg2 as getConnectionCallbackType,
             format
           );
           return originalGetConnection.call(pool, arg1, patchFn);
         }
         if (arguments.length === 3 && typeof arg3 === 'function') {
           const patchFn = thisPlugin._getConnectionCallbackPatchFn(
-            arg3,
+            arg3 as getConnectionCallbackType,
             format
           );
           return originalGetConnection.call(pool, arg1, arg2, patchFn);
@@ -207,25 +209,23 @@ export class MySQLInstrumentation extends InstrumentationBase<
     };
   }
 
-  private _getConnectionCallbackPatchFn(cb: Function, format: formatType) {
+  private _getConnectionCallbackPatchFn(cb: getConnectionCallbackType, format: formatType) {
     const thisPlugin = this;
     const activeContext = context.active();
-    return function () {
-      if (arguments[1]) {
+    return function (this: any, err: mysqlTypes.MysqlError, connection: mysqlTypes.PoolConnection) {
+      if (connection) {
         // this is the callback passed into a query
         // no need to unwrap
-        if (!isWrapped(arguments[1].query)) {
+        if (!isWrapped(connection.query)) {
           thisPlugin._wrap(
-            arguments[1],
+            connection,
             'query',
-            thisPlugin._patchQuery(arguments[1], format)
+            thisPlugin._patchQuery(connection, format)
           );
         }
       }
       if (typeof cb === 'function') {
-        context.with(activeContext, () => {
-          cb(...arguments);
-        });
+        context.with(activeContext, cb, this, err, connection);
       }
     };
   }

diff --git a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
@@ -438,13 +438,8 @@ describe('mysql@2.x', () => {
             assert.ifError(err);
             assert.ok(res);
             assert.strictEqual(res[0].solution, 1);
-            const spans = memoryExporter.getFinishedSpans();
-            assert.strictEqual(spans.length, 1);
-            const querySpan = spans[0];
-            assert.strictEqual(
-              querySpan.parentSpanId,
-              parentSpan.spanContext().spanId
-            );
+            const actualSpan = trace.getSpan(context.active());
+            assert.strictEqual(actualSpan, parentSpan);
             done();
           });
         });
@@ -636,13 +631,8 @@ describe('mysql@2.x', () => {
             assert.ifError(err);
             assert.ok(res);
             assert.strictEqual(res[0].solution, 1);
-            const spans = memoryExporter.getFinishedSpans();
-            assert.strictEqual(spans.length, 1);
-            const querySpan = spans[0];
-            assert.strictEqual(
-              querySpan.parentSpanId,
-              parentSpan.spanContext().spanId
-            );
+            const actualSpan = trace.getSpan(context.active());
+            assert.strictEqual(actualSpan, parentSpan);
             done();
           });
         });