fix(mysql): set proper context for outgoing queries

plugins/node/opentelemetry-instrumentation-mysql/package.json

@@ -51,11 +51,13 @@
     "@opentelemetry/sdk-trace-base": "^1.3.1",
     "@types/mocha": "7.0.2",
     "@types/node": "16.11.21",
+    "@types/sinon":  "10.0.13",
     "gts": "3.1.0",
     "mocha": "7.2.0",
     "mysql": "2.18.1",
     "nyc": "15.1.0",
     "rimraf": "3.0.2",
+    "sinon": "14.0.0",
     "ts-mocha": "10.0.0",
     "typescript": "4.3.5"
   },

plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts

@@ -16,7 +16,9 @@
 
 import {
   context,
+  Context,
   diag,
+  trace,
   Span,
   SpanKind,
   SpanStatusCode,
@@ -285,11 +287,16 @@ export class MySQLInstrumentation extends InstrumentationBase<
           arg => typeof arg === 'function'
         );
 
+        const parentContext = context.active();
+
         if (cbIndex === -1) {
-          const streamableQuery: mysqlTypes.Query = originalQuery.apply(
-            connection,
-            arguments
+          const streamableQuery: mysqlTypes.Query = context.with(
+            trace.setSpan(context.active(), span),
+            () => {
+              return originalQuery.apply(connection, arguments);
+            }
           );
+          context.bind(parentContext, streamableQuery);
 
           return streamableQuery
             .on('error', err =>
@@ -305,16 +312,18 @@ export class MySQLInstrumentation extends InstrumentationBase<
           thisPlugin._wrap(
             arguments,
             cbIndex,
-            thisPlugin._patchCallbackQuery(span)
+            thisPlugin._patchCallbackQuery(span, parentContext)
           );
 
-          return originalQuery.apply(connection, arguments);
+          return context.with(trace.setSpan(context.active(), span), () => {
+            return originalQuery.apply(connection, arguments);
+          });
         }
       };
     };
   }
 
-  private _patchCallbackQuery(span: Span) {
+  private _patchCallbackQuery(span: Span, parentContext: Context) {
     return (originalCallback: Function) => {
       return function (
         err: mysqlTypes.MysqlError | null,
@@ -328,7 +337,9 @@ export class MySQLInstrumentation extends InstrumentationBase<
           });
         }
         span.end();
-        return originalCallback(...arguments);
+        return context.with(parentContext, () =>
+          originalCallback(...arguments)
+        );
       };
     };
   }

plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { context, trace, SpanStatusCode } from '@opentelemetry/api';
+import { context, Context, trace, SpanStatusCode } from '@opentelemetry/api';
 import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 import {
   DbSystemValues,
@@ -29,6 +29,7 @@ import {
 } from '@opentelemetry/sdk-trace-base';
 import * as assert from 'assert';
 import { MySQLInstrumentation } from '../src';
+import * as sinon from 'sinon';
 
 const port = Number(process.env.MYSQL_PORT) || 33306;
 const database = process.env.MYSQL_DATABASE || 'test_db';
@@ -291,6 +292,70 @@ describe('mysql@2.x', () => {
         });
       });
     });
+
+    describe('active span context', () => {
+      afterEach(() => {
+        sinon.restore();
+      });
+
+      function assertParent(parentContext: Context) {
+        const anyConn = connection as any;
+        const originalImplyConnect = anyConn._implyConnect.bind(connection);
+        return sinon.stub(anyConn, '_implyConnect').callsFake(() => {
+          const activeSpan = trace.getSpan(
+            context.active()
+          ) as unknown as ReadableSpan;
+          const parentSpanContext = trace.getSpanContext(parentContext);
+          assert.strictEqual(
+            activeSpan.spanContext().traceId,
+            parentSpanContext?.traceId
+          );
+          assert.strictEqual(
+            activeSpan.parentSpanId,
+            parentSpanContext?.spanId
+          );
+          assert.notStrictEqual(
+            activeSpan.spanContext().spanId,
+            parentSpanContext?.spanId
+          );
+          originalImplyConnect();
+        });
+      }
+
+      it('should have proper context active for connection.query(text: string)', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const parentContext = context.active();
+          const stub = assertParent(parentContext);
+
+          const query = connection.query('select 1');
+
+          query.on('result', () => {
+            assert.strictEqual(context.active(), parentContext);
+          });
+
+          query.on('end', () => {
+            assert.strictEqual(context.active(), parentContext);
+            sinon.assert.called(stub);
+            done();
+          });
+        });
+      });
+
+      it('should have proper context active for connection.query(text: string, callback)', done => {
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const parentContext = context.active();
+          const stub = assertParent(parentContext);
+
+          connection.query('select 1', () => {
+            sinon.assert.called(stub);
+            assert.strictEqual(context.active(), parentContext);
+            done();
+          });
+        });
+      });
+    });
   });
 
   describe('#Pool', () => {