From f2faeda2993536547fecbbec63a4bc7ad83583e6 Mon Sep 17 00:00:00 2001 From: Joel Mukuthu Date: Mon, 3 Apr 2023 10:46:59 +0200 Subject: [PATCH 1/2] feat(instrumentation-pg): support custom SQL Commenter attributes --- .../src/utils.ts | 15 ++++-- .../test/utils.test.ts | 54 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts index 2e4176817d..883086b451 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts @@ -47,6 +47,8 @@ function arrayStringifyHelper(arr: Array): string { return '[' + arr.toString() + ']'; } +type SqlCommenterAttributes = Record; + /** * Helper function to get a low cardinality span name from whatever info we have * about the query. @@ -284,7 +286,11 @@ function fixedEncodeURIComponent(str: string) { ); } -export function addSqlCommenterComment(span: Span, query: string): string { +export function addSqlCommenterComment( + span: Span, + query: string, + customAttributes: SqlCommenterAttributes = {} +): string { if (typeof query !== 'string' || query.length === 0) { return query; } @@ -296,7 +302,7 @@ export function addSqlCommenterComment(span: Span, query: string): string { } const propagator = new W3CTraceContextPropagator(); - const headers: { [key: string]: string } = {}; + const headers: SqlCommenterAttributes = customAttributes; propagator.inject( trace.setSpan(ROOT_CONTEXT, span), headers, @@ -311,9 +317,10 @@ export function addSqlCommenterComment(span: Span, query: string): string { } const commentString = sortedKeys - .map(key => { + .map((key) => { + const encodedKey = fixedEncodeURIComponent(key); const encodedValue = fixedEncodeURIComponent(headers[key]); - return `${key}='${encodedValue}'`; + return `${encodedKey}='${encodedValue}'`; }) .join(','); diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts index 2e6c8ca0d7..cf657ff415 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts @@ -289,5 +289,59 @@ describe('utils.ts', () => { "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01',tracestate='foo%3D%27bar%2Cbaz%3D%27qux%21%28%29%2A%27%2Chack%3D%27DROP%20TABLE'*/" ); }); + + it('supports custom attributes', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment( + trace.wrapSpanContext(spanContext), + query, + { controller: 'foo', action: 'bar' } + ), + "SELECT * from FOO; /*action='bar',controller='foo',traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/" + ); + }); + + it('escapes special characters in keys', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment( + trace.wrapSpanContext(spanContext), + query, + { "'DROP TABLE": 'foo' } + ), + "SELECT * from FOO; /*%27DROP%20TABLE='foo',traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/" + ); + }); + + it('overwrites custom attributes with trace context keys', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment( + trace.wrapSpanContext(spanContext), + query, + { traceparent: 'foo' } + ), + "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/" + ); + }); }); }); From be38c3e4a401d38205526652fed7d5ce93512e12 Mon Sep 17 00:00:00 2001 From: Joel Mukuthu Date: Mon, 3 Apr 2023 10:54:18 +0200 Subject: [PATCH 2/2] feat(instrumentation-pg): get custom SQL Commenter attributes from active context --- .../src/instrumentation.ts | 16 +++++++++++++--- .../src/utils.ts | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index a5eff14498..622b726200 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts @@ -215,18 +215,28 @@ export class PgInstrumentation extends InstrumentationBase { // Modify query text w/ a tracing comment before invoking original for // tracing, but only if args[0] has one of our expected shapes. // - // TODO: remove the `as ...` casts below when the TS version is upgraded. + // TODO: remove the `as ...` casts on `arg0` below when the TS version is upgraded. // Newer TS versions will use the result of firstArgIsQueryObjectWithText // to properly narrow arg0, but TS 4.3.5 does not. if (instrumentationConfig.addSqlCommenterCommentToQueries) { + const customSqlCommenterAttributes = context + .active() + .getValue( + Symbol('customSqlCommenterAttributes') + ) as utils.SqlCommenterAttributes; args[0] = firstArgIsString - ? utils.addSqlCommenterComment(span, arg0 as string) + ? utils.addSqlCommenterComment( + span, + arg0 as string, + customSqlCommenterAttributes + ) : firstArgIsQueryObjectWithText ? { ...(arg0 as utils.ObjectWithText), text: utils.addSqlCommenterComment( span, - (arg0 as utils.ObjectWithText).text + (arg0 as utils.ObjectWithText).text, + customSqlCommenterAttributes ), } : args[0]; diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts index 883086b451..c91c8ae75e 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts @@ -47,7 +47,7 @@ function arrayStringifyHelper(arr: Array): string { return '[' + arr.toString() + ']'; } -type SqlCommenterAttributes = Record; +export type SqlCommenterAttributes = Record; /** * Helper function to get a low cardinality span name from whatever info we have