fix(mysql): bind get connection callback to active context

Signed-off-by: Simon Stone <simonstone1@gmail.com>
open-telemetry · Jul 6, 2021 · 3ebf81d · 3ebf81d
1 parent 52b46aa
commit 3ebf81d
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 2 deletions.
diff --git 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
@@ -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();
+          });
+        });
+      });
+    });
   });
 });