From 3ebf81d0f5ed24479e5df5302a0fc309bbbc865f Mon Sep 17 00:00:00 2001
From: Simon Stone <simonstone1@gmail.com>
Date: Mon, 5 Jul 2021 16:25:56 +0100
Subject: [PATCH] fix(mysql): bind get connection callback to active context

Signed-off-by: Simon Stone <simonstone1@gmail.com>
---
 .../src/instrumentation.ts                    | 13 ++++-
 .../test/index.test.ts                        | 48 +++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
index 301b85e51d..dec0427606 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
@@ -14,7 +14,13 @@
  * limitations under the License.
  */
 
-import { diag, Span, SpanKind, SpanStatusCode } from '@opentelemetry/api';
+import {
+  context,
+  diag,
+  Span,
+  SpanKind,
+  SpanStatusCode,
+} from '@opentelemetry/api';
 import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
@@ -203,6 +209,7 @@ export class MySQLInstrumentation extends InstrumentationBase<
 
   private _getConnectionCallbackPatchFn(cb: Function, format: formatType) {
     const thisPlugin = this;
+    const activeContext = context.active();
     return function () {
       if (arguments[1]) {
         // this is the callback passed into a query
@@ -216,7 +223,9 @@ export class MySQLInstrumentation extends InstrumentationBase<
         }
       }
       if (typeof cb === 'function') {
-        cb(...arguments);
+        context.with(activeContext, () => {
+          cb(...arguments);
+        });
       }
     };
   }
diff --git a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
index fbf726281b..afd468f17c 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
@@ -426,6 +426,30 @@ describe('mysql@2.x', () => {
         });
       });
     });
+
+    it('should propagate active context to callback', done => {
+      const parentSpan = provider.getTracer('default').startSpan('test span');
+      context.with(trace.setSpan(context.active(), parentSpan), () => {
+        pool.getConnection((err, connection) => {
+          assert.ifError(err);
+          assert.ok(connection);
+          const sql = 'SELECT ? as solution';
+          connection.query(sql, 1, (err, res) => {
+            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
+            );
+            done();
+          });
+        });
+      });
+    });
   });
 
   describe('#PoolCluster', () => {
@@ -600,6 +624,30 @@ describe('mysql@2.x', () => {
         }
       );
     });
+
+    it('should propagate active context to callback', done => {
+      const parentSpan = provider.getTracer('default').startSpan('test span');
+      context.with(trace.setSpan(context.active(), parentSpan), () => {
+        poolCluster.getConnection((err, connection) => {
+          assert.ifError(err);
+          assert.ok(connection);
+          const sql = 'SELECT ? as solution';
+          connection.query(sql, 1, (err, res) => {
+            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
+            );
+            done();
+          });
+        });
+      });
+    });
   });
 });