diff --git a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts index 82ea9770d4..8a0652d503 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, @@ -28,6 +34,11 @@ 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 > { @@ -176,21 +187,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); @@ -201,22 +212,30 @@ export class MySQLInstrumentation extends InstrumentationBase< }; } - private _getConnectionCallbackPatchFn(cb: Function, format: formatType) { + private _getConnectionCallbackPatchFn( + cb: getConnectionCallbackType, + format: formatType + ) { const thisPlugin = this; - return function () { - if (arguments[1]) { + const activeContext = context.active(); + 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') { - 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 index fbf726281b..060b56e0cf 100644 --- a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts +++ b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts @@ -426,6 +426,25 @@ 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 actualSpan = trace.getSpan(context.active()); + assert.strictEqual(actualSpan, parentSpan); + done(); + }); + }); + }); + }); }); describe('#PoolCluster', () => { @@ -600,6 +619,25 @@ 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 actualSpan = trace.getSpan(context.active()); + assert.strictEqual(actualSpan, parentSpan); + done(); + }); + }); + }); + }); }); });