diff --git a/plugins/node/opentelemetry-instrumentation-mysql/package.json b/plugins/node/opentelemetry-instrumentation-mysql/package.json index fa6bdae066..4a63a19da5 100644 --- a/plugins/node/opentelemetry-instrumentation-mysql/package.json +++ b/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" }, diff --git a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts index 795e55b9f7..66809f9688 100644 --- a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts +++ b/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) + ); }; }; } diff --git a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts index 3f3c9c808f..e5e5804336 100644 --- a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts +++ b/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', () => {