From a0003e76fc46afbbee2558a7d21906be7c9cb1d1 Mon Sep 17 00:00:00 2001 From: Henri Normak Date: Thu, 24 Nov 2022 17:51:42 +0200 Subject: [PATCH] feat: add sqlcommenter comment with trace context to queries in pg instrumentation (#1286) --- .../README.md | 2 + .../package.json | 3 + .../src/instrumentation.ts | 16 +++ .../src/types.ts | 6 + .../src/utils.ts | 69 ++++++++++ .../test/pg.test.ts | 123 ++++++++++++++++++ .../test/utils.test.ts | 99 +++++++++++++- 7 files changed, 317 insertions(+), 1 deletion(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/README.md b/plugins/node/opentelemetry-instrumentation-pg/README.md index 578b2cc173..b14e036343 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/README.md +++ b/plugins/node/opentelemetry-instrumentation-pg/README.md @@ -48,6 +48,8 @@ PostgreSQL instrumentation has few options available to choose from. You can set | ------- | ---- | ----------- | | [`enhancedDatabaseReporting`](./src/types.ts#L30) | `boolean` | If true, additional information about query parameters and results will be attached (as `attributes`) to spans representing database operations | | `responseHook` | `PgInstrumentationExecutionResponseHook` (function) | Function for adding custom attributes from db response | +| `requireParentSpan` | `boolean` | If true, requires a parent span to create new spans (default false) | +| `addSqlCommenterCommentToQueries` | `boolean` | If true, adds [sqlcommenter](https://github.com/open-telemetry/opentelemetry-sqlcommenter) specification compliant comment to queries with tracing context (default false). _NOTE: A comment will not be added to queries that already contain `--` or `/* ... */` in them, even if these are not actually part of comments_ | ## Useful links diff --git a/plugins/node/opentelemetry-instrumentation-pg/package.json b/plugins/node/opentelemetry-instrumentation-pg/package.json index 917139df6e..fdf69fe8be 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/package.json +++ b/plugins/node/opentelemetry-instrumentation-pg/package.json @@ -60,6 +60,7 @@ "@opentelemetry/sdk-trace-node": "^1.8.0", "@types/mocha": "7.0.2", "@types/node": "18.11.7", + "@types/sinon": "10.0.2", "cross-env": "7.0.3", "gts": "3.1.0", "mocha": "7.2.0", @@ -67,11 +68,13 @@ "pg": "8.7.1", "pg-pool": "3.4.1", "rimraf": "3.0.2", + "sinon": "14.0.0", "test-all-versions": "5.0.1", "ts-mocha": "10.0.0", "typescript": "4.3.5" }, "dependencies": { + "@opentelemetry/core": "^1.8.0", "@opentelemetry/instrumentation": "^0.34.0", "@opentelemetry/semantic-conventions": "^1.0.0", "@types/pg": "8.6.1", diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index 7c1b6d77f2..f309555951 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts @@ -178,6 +178,7 @@ export class PgInstrumentation extends InstrumentationBase { // Handle different client.query(...) signatures if (typeof args[0] === 'string') { const query = args[0]; + if (args.length > 1 && args[1] instanceof Array) { const params = args[1]; span = utils.handleParameterizedQuery.call( @@ -195,14 +196,29 @@ export class PgInstrumentation extends InstrumentationBase { query ); } + + if (plugin.getConfig().addSqlCommenterCommentToQueries) { + // Modify the query with a tracing comment + args[0] = utils.addSqlCommenterComment(span, args[0]); + } } else if (typeof args[0] === 'object') { const queryConfig = args[0] as NormalizedQueryConfig; + span = utils.handleConfigQuery.call( this, plugin.tracer, plugin.getConfig(), queryConfig ); + + if (plugin.getConfig().addSqlCommenterCommentToQueries) { + // Copy the query config instead of writing to args[0].text so that we don't modify user's + // original query config reference + args[0] = { + ...queryConfig, + text: utils.addSqlCommenterComment(span, queryConfig.text), + }; + } } else { return utils.handleInvalidQuery.call( this, diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/types.ts b/plugins/node/opentelemetry-instrumentation-pg/src/types.ts index d37ab0323f..3e0d280cad 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/types.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/types.ts @@ -46,4 +46,10 @@ export interface PgInstrumentationConfig extends InstrumentationConfig { * @default false */ requireParentSpan?: boolean; + + /** + * If true, queries are modified to also include a comment with + * the tracing context, following the {@link https://github.com/open-telemetry/opentelemetry-sqlcommenter sqlcommenter} format + */ + addSqlCommenterCommentToQueries?: boolean; } diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts index 7e5359c4ec..048c0d7ad0 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts @@ -24,7 +24,10 @@ import { diag, INVALID_SPAN_CONTEXT, Attributes, + defaultTextMapSetter, + ROOT_CONTEXT, } from '@opentelemetry/api'; +import { W3CTraceContextPropagator } from '@opentelemetry/core'; import { AttributeNames } from './enums/AttributeNames'; import { SemanticAttributes, @@ -289,3 +292,69 @@ export function patchClientConnectCallback( cb.call(this, err); }; } + +// NOTE: This function currently is returning false-positives +// in cases where comment characters appear in string literals +// ("SELECT '-- not a comment';" would return true, although has no comment) +function hasValidSqlComment(query: string): boolean { + const indexOpeningDashDashComment = query.indexOf('--'); + if (indexOpeningDashDashComment >= 0) { + return true; + } + + const indexOpeningSlashComment = query.indexOf('/*'); + if (indexOpeningSlashComment < 0) { + return false; + } + + const indexClosingSlashComment = query.indexOf('*/'); + return indexOpeningDashDashComment < indexClosingSlashComment; +} + +// sqlcommenter specification (https://google.github.io/sqlcommenter/spec/#value-serialization) +// expects us to URL encode based on the RFC 3986 spec (https://en.wikipedia.org/wiki/Percent-encoding), +// but encodeURIComponent does not handle some characters correctly (! ' ( ) *), +// which means we need special handling for this +// https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/encodeURIComponent +function fixedEncodeURIComponent(str: string) { + return encodeURIComponent(str).replace( + /[!'()*]/g, + c => `%${c.charCodeAt(0).toString(16).toUpperCase()}` + ); +} + +export function addSqlCommenterComment(span: Span, query: string): string { + if (typeof query !== 'string' || query.length === 0) { + return query; + } + + // As per sqlcommenter spec we shall not add a comment if there already is a comment + // in the query + if (hasValidSqlComment(query)) { + return query; + } + + const propagator = new W3CTraceContextPropagator(); + const headers: { [key: string]: string } = {}; + propagator.inject( + trace.setSpan(ROOT_CONTEXT, span), + headers, + defaultTextMapSetter + ); + + // sqlcommenter spec requires keys in the comment to be sorted lexicographically + const sortedKeys = Object.keys(headers).sort(); + + if (sortedKeys.length === 0) { + return query; + } + + const commentString = sortedKeys + .map(key => { + const encodedValue = fixedEncodeURIComponent(headers[key]); + return `${key}='${encodedValue}'`; + }) + .join(','); + + return `${query} /*${commentString}*/`; +} diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts index d43eb7537f..e91b0034fb 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts @@ -32,6 +32,7 @@ import { } from '@opentelemetry/sdk-trace-base'; import * as assert from 'assert'; import type * as pg from 'pg'; +import * as sinon from 'sinon'; import { PgInstrumentation, PgInstrumentationConfig, @@ -44,6 +45,7 @@ import { DbSystemValues, } from '@opentelemetry/semantic-conventions'; import { isSupported } from './utils'; +import { addSqlCommenterComment } from '../src/utils'; const pgVersion = require('pg/package.json').version; const nodeVersion = process.versions.node; @@ -110,6 +112,12 @@ describe('pg', () => { const testPostgresLocally = process.env.RUN_POSTGRES_TESTS_LOCAL; // For local: spins up local postgres db via docker const shouldTest = testPostgres || testPostgresLocally; // Skips these tests if false (default) + function getExecutedQueries() { + return (client as any).queryQueue.push.args.flat() as (pg.Query & { + text?: string; + })[]; + } + before(async function () { const skipForUnsupported = process.env.IN_TAV && !isSupported(nodeVersion, pgVersion); @@ -156,11 +164,16 @@ describe('pg', () => { beforeEach(() => { contextManager = new AsyncHooksContextManager().enable(); context.setGlobalContextManager(contextManager); + + // Add a spy on the underlying client's internal query queue so that + // we could assert on what the final queries are that are executed + sinon.spy((client as any).queryQueue, 'push'); }); afterEach(() => { memoryExporter.reset(); context.disable(); + sinon.restore(); }); it('should return an instrumentation', () => { @@ -649,6 +662,116 @@ describe('pg', () => { }); }); + it('should not add sqlcommenter comment when flag is not specified', async () => { + const span = tracer.startSpan('test span'); + await context.with(trace.setSpan(context.active(), span), async () => { + try { + const query = 'SELECT NOW()'; + const resPromise = await client.query(query); + assert.ok(resPromise); + + const [span] = memoryExporter.getFinishedSpans(); + assert.ok(span); + + const commentedQuery = addSqlCommenterComment( + trace.wrapSpanContext(span.spanContext()), + query + ); + + const executedQueries = getExecutedQueries(); + assert.equal(executedQueries.length, 1); + assert.equal(executedQueries[0].text, query); + assert.notEqual(query, commentedQuery); + } catch (e) { + assert.ok(false, e.message); + } + }); + }); + + it('should not add sqlcommenter comment with client.query({text, callback}) when flag is not specified', done => { + const span = tracer.startSpan('test span'); + context.with(trace.setSpan(context.active(), span), () => { + const query = 'SELECT NOW()'; + client.query({ + text: query, + callback: (err: Error, res: pg.QueryResult) => { + assert.strictEqual(err, null); + assert.ok(res); + + const [span] = memoryExporter.getFinishedSpans(); + const commentedQuery = addSqlCommenterComment( + trace.wrapSpanContext(span.spanContext()), + query + ); + + const executedQueries = getExecutedQueries(); + assert.equal(executedQueries.length, 1); + assert.equal(executedQueries[0].text, query); + assert.notEqual(query, commentedQuery); + done(); + }, + } as pg.QueryConfig); + }); + }); + + it('should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified', async () => { + instrumentation.setConfig({ + addSqlCommenterCommentToQueries: true, + }); + + const span = tracer.startSpan('test span'); + await context.with(trace.setSpan(context.active(), span), async () => { + try { + const query = 'SELECT NOW()'; + const resPromise = await client.query(query); + assert.ok(resPromise); + + const [span] = memoryExporter.getFinishedSpans(); + const commentedQuery = addSqlCommenterComment( + trace.wrapSpanContext(span.spanContext()), + query + ); + + const executedQueries = getExecutedQueries(); + assert.equal(executedQueries.length, 1); + assert.equal(executedQueries[0].text, commentedQuery); + assert.notEqual(query, commentedQuery); + } catch (e) { + assert.ok(false, e.message); + } + }); + }); + + it('should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified with client.query({text, callback})', done => { + instrumentation.setConfig({ + addSqlCommenterCommentToQueries: true, + }); + + const span = tracer.startSpan('test span'); + context.with(trace.setSpan(context.active(), span), () => { + const query = 'SELECT NOW()'; + client.query({ + text: query, + callback: (err: Error, res: pg.QueryResult) => { + assert.strictEqual(err, null); + assert.ok(res); + + const [span] = memoryExporter.getFinishedSpans(); + const commentedQuery = addSqlCommenterComment( + trace.wrapSpanContext(span.spanContext()), + query + ); + + const executedQueries = getExecutedQueries(); + assert.equal(executedQueries.length, 1); + assert.equal(executedQueries[0].text, commentedQuery); + assert.notEqual(query, commentedQuery); + done(); + }, + } as pg.QueryConfig); + }); + }); + it('should not generate traces for client.query() when requireParentSpan=true is specified', done => { instrumentation.setConfig({ requireParentSpan: true, diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts index 6bb03076bf..4b43e4da95 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts @@ -14,7 +14,14 @@ * limitations under the License. */ -import { context, INVALID_SPAN_CONTEXT, trace } from '@opentelemetry/api'; +import { + context, + createTraceState, + INVALID_SPAN_CONTEXT, + SpanContext, + trace, + TraceFlags, +} from '@opentelemetry/api'; import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks'; import { InstrumentationConfig } from '@opentelemetry/instrumentation'; import { @@ -204,4 +211,94 @@ describe('utils.ts', () => { assert.strictEqual(pgValues, '[0]'); }); }); + + describe('addSqlCommenterComment', () => { + it('adds comment to a simple query', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment(trace.wrapSpanContext(spanContext), query), + "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/" + ); + }); + + it('does not add a comment if query already has a comment', () => { + const span = trace.wrapSpanContext({ + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }); + + const blockComment = 'SELECT * from FOO; /* Test comment */'; + assert.strictEqual( + utils.addSqlCommenterComment(span, blockComment), + blockComment + ); + + const dashedComment = 'SELECT * from FOO; -- Test comment'; + assert.strictEqual( + utils.addSqlCommenterComment(span, dashedComment), + dashedComment + ); + }); + + it('does not add a comment to an empty query', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + }; + + assert.strictEqual( + utils.addSqlCommenterComment(trace.wrapSpanContext(spanContext), ''), + '' + ); + }); + + it('does not add a comment if span context is invalid', () => { + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment( + trace.wrapSpanContext(INVALID_SPAN_CONTEXT), + query + ), + query + ); + }); + + it('correctly also sets trace state', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState('foo=bar,baz=qux'), + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment(trace.wrapSpanContext(spanContext), query), + "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01',tracestate='foo%3Dbar%2Cbaz%3Dqux'*/" + ); + }); + + it('escapes special characters in values', () => { + const spanContext: SpanContext = { + traceId: 'd4cda95b652f4a1592b449d5929fda1b', + spanId: '6e0c63257de34c92', + traceFlags: TraceFlags.SAMPLED, + traceState: createTraceState("foo='bar,baz='qux!()*',hack='DROP TABLE"), + }; + + const query = 'SELECT * from FOO;'; + assert.strictEqual( + utils.addSqlCommenterComment(trace.wrapSpanContext(spanContext), query), + "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01',tracestate='foo%3D%27bar%2Cbaz%3D%27qux%21%28%29%2A%27%2Chack%3D%27DROP%20TABLE'*/" + ); + }); + }); });