From d6f8af008163a72c8b590ae002a7c780aab5b866 Mon Sep 17 00:00:00 2001 From: Alon Katz Date: Tue, 30 Nov 2021 15:32:33 +0200 Subject: [PATCH 1/3] add post-query custom attributes hook for spans --- .../README.md | 1 + .../src/instrumentation.ts | 22 ++-- .../src/types.ts | 16 +++ .../src/utils.ts | 20 ++++ .../test/pg.test.ts | 106 ++++++++++++++++++ 5 files changed, 156 insertions(+), 9 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/README.md b/plugins/node/opentelemetry-instrumentation-pg/README.md index f99c4574cb..5d37961225 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/README.md +++ b/plugins/node/opentelemetry-instrumentation-pg/README.md @@ -48,6 +48,7 @@ PostgreSQL instrumentation has few options available to choose from. You can set | Options | Type | Description | | ------- | ---- | ----------- | | [`enhancedDatabaseReporting`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/opentelemetry-instrumentation-pg/src/pg.ts#L48) | `boolean` | If true, additional information about query parameters and results will be attached (as `attributes`) to spans representing database operations | +| `postQueryHook` | `PgPostQueryHookFunction` (function) | Function for adding custom attributes before the query is resolved | | `responseHook` | `PgInstrumentationExecutionResponseHook` (function) | Function for adding custom attributes from db response | ## Supported Versions diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index 7260e4fb2c..0906ad2c6f 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts @@ -36,6 +36,7 @@ import { PgPoolExtended, PgPoolCallback, PgInstrumentationConfig, + QueryContext, } from './types'; import * as utils from './utils'; import { AttributeNames } from './enums/AttributeNames'; @@ -114,7 +115,9 @@ export class PgInstrumentation extends InstrumentationBase { `Patching ${PgInstrumentation.COMPONENT}.Client.prototype.query` ); return function query(this: PgClientExtended, ...args: unknown[]) { + const pluginConfig = plugin.getConfig() as PgInstrumentationConfig; let span: Span; + let postQueryHookParams: QueryContext; // Handle different client.query(...) signatures if (typeof args[0] === 'string') { @@ -124,21 +127,24 @@ export class PgInstrumentation extends InstrumentationBase { span = utils.handleParameterizedQuery.call( this, plugin.tracer, - plugin.getConfig() as PgInstrumentationConfig, + pluginConfig, query, params ); + postQueryHookParams = { span, query, params }; } else { span = utils.handleTextQuery.call(this, plugin.tracer, query); + postQueryHookParams = { span, query }; } } else if (typeof args[0] === 'object') { const queryConfig = args[0] as NormalizedQueryConfig; span = utils.handleConfigQuery.call( this, plugin.tracer, - plugin.getConfig() as PgInstrumentationConfig, + pluginConfig, queryConfig ); + postQueryHookParams = { span, config: queryConfig }; } else { return utils.handleInvalidQuery.call( this, @@ -154,7 +160,7 @@ export class PgInstrumentation extends InstrumentationBase { if (typeof args[args.length - 1] === 'function') { // Patch ParameterQuery callback args[args.length - 1] = utils.patchCallback( - plugin.getConfig() as PgInstrumentationConfig, + pluginConfig, span, args[args.length - 1] as PostgresCallback ); @@ -170,7 +176,7 @@ export class PgInstrumentation extends InstrumentationBase { ) { // Patch ConfigQuery callback let callback = utils.patchCallback( - plugin.getConfig() as PgInstrumentationConfig, + pluginConfig, span, (args[0] as NormalizedQueryConfig).callback! ); @@ -185,6 +191,8 @@ export class PgInstrumentation extends InstrumentationBase { } } + utils.handlePostQueryHook(pluginConfig, postQueryHookParams); + // Perform the original query const result: unknown = original.apply(this, args as never); @@ -194,11 +202,7 @@ export class PgInstrumentation extends InstrumentationBase { .then((result: unknown) => { // Return a pass-along promise which ends the span and then goes to user's orig resolvers return new Promise(resolve => { - utils.handleExecutionResult( - plugin.getConfig() as PgInstrumentationConfig, - span, - result - ); + utils.handleExecutionResult(pluginConfig, span, result); span.end(); resolve(result); }); diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/types.ts b/plugins/node/opentelemetry-instrumentation-pg/src/types.ts index 8728dc1028..acb83c8e42 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/types.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/types.ts @@ -27,12 +27,28 @@ export interface PgInstrumentationExecutionResponseHook { (span: api.Span, responseInfo: PgResponseHookInformation): void; } +export interface QueryContext { + span: api.Span; + query?: string; + config?: NormalizedQueryConfig; + params?: unknown[]; +} + +export interface PgPostQueryHookFunction { + (ctx: QueryContext): void; +} + export interface PgInstrumentationConfig extends InstrumentationConfig { /** * If true, additional information about query parameters will be attached (as `attributes`) to spans representing */ enhancedDatabaseReporting?: boolean; + /** + * Function for adding custom attributes before the query is resolved + */ + postQueryHook?: PgPostQueryHookFunction; + /** * Hook that allows adding custom span attributes based on the data * returned from "query" Pg actions. diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts index 87efaf6400..419a66c401 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts @@ -34,6 +34,7 @@ import { PgPoolCallback, PgPoolExtended, PgInstrumentationConfig, + QueryContext, } from './types'; import * as pgTypes from 'pg'; import { PgInstrumentation } from './'; @@ -169,6 +170,25 @@ export function handleInvalidQuery( return result; } +export function handlePostQueryHook( + config: PgInstrumentationConfig, + ctx: QueryContext +) { + if (typeof config.postQueryHook === 'function') { + safeExecuteInTheMiddle( + () => { + config.postQueryHook!(ctx); + }, + err => { + if (err) { + diag.error('Error running post query hook', err); + } + }, + true + ); + } +} + export function handleExecutionResult( config: PgInstrumentationConfig, span: Span, diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts index 2c26cf4c8f..7fb7daf451 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts @@ -563,5 +563,111 @@ describe('pg', () => { client.query('SELECT NOW()').then(queryHandler); }); }); + + it('should call postQueryHook with query text if set', async () => { + instrumentation.disable(); + let called = false; + const query = 'SELECT NOW()'; + const config: PgInstrumentationConfig = { + postQueryHook: ctx => { + called = true; + assert.strictEqual(ctx.query, query); + assert.strictEqual(ctx.params, undefined); + }, + }; + instrumentation.setConfig(config); + instrumentation.enable(); + + const attributes = { + ...DEFAULT_ATTRIBUTES, + [SemanticAttributes.DB_STATEMENT]: query, + }; + const events: TimedEvent[] = []; + const span = tracer.startSpan('test span'); + await context.with(trace.setSpan(context.active(), span), async () => { + try { + const resPromise = await client.query(query); + assert.ok(resPromise); + runCallbackTest(span, attributes, events); + } catch (e) { + assert.ok(false, e.message); + } + }); + assert.strictEqual(called, true); + }); + it('should call postQueryHook with query text and params if set', async () => { + instrumentation.disable(); + let called = false; + const values = ['0']; + const query = 'SELECT $1::text'; + const config: PgInstrumentationConfig = { + postQueryHook: ctx => { + called = true; + assert.strictEqual(ctx.query, query); + assert.strictEqual(ctx.params, values); + }, + }; + instrumentation.setConfig(config); + instrumentation.enable(); + + const attributes = { + ...DEFAULT_ATTRIBUTES, + [SemanticAttributes.DB_STATEMENT]: query, + }; + const events: TimedEvent[] = []; + const span = tracer.startSpan('test span'); + await context.with(trace.setSpan(context.active(), span), async () => { + const resPromise = await client.query(query, values); + try { + assert.ok(resPromise); + runCallbackTest(span, attributes, events); + } catch (e) { + assert.ok(false, e.message); + } + }); + assert.strictEqual(called, true); + }); + it('should call postQueryHook with query config if set', async () => { + instrumentation.disable(); + const name = 'fetch-text'; + const query = 'SELECT $1::text'; + const values = ['0']; + let called = false; + const config: PgInstrumentationConfig = { + postQueryHook: ctx => { + called = true; + if (!ctx.config) { + assert.ok(false, 'ctx.config was undefined'); + } + assert.strictEqual(ctx.config.text, query); + assert.strictEqual(ctx.config.values, values); + }, + }; + instrumentation.setConfig(config); + instrumentation.enable(); + + const attributes = { + ...DEFAULT_ATTRIBUTES, + [AttributeNames.PG_PLAN]: name, + [SemanticAttributes.DB_STATEMENT]: query, + }; + const events: TimedEvent[] = []; + const span = tracer.startSpan('test span'); + + await context.with(trace.setSpan(context.active(), span), async () => { + try { + const resPromise = await client.query({ + name: name, + text: query, + values: values, + }); + assert.strictEqual(resPromise.command, 'SELECT'); + runCallbackTest(span, attributes, events); + } catch (e) { + assert.ok(false, e.message); + } + }); + assert.strictEqual(called, true); + }); }); }); From 85d561152a84e0e63b094f64ae7ab958130ba63b Mon Sep 17 00:00:00 2001 From: Alon Katz Date: Tue, 30 Nov 2021 16:44:47 +0200 Subject: [PATCH 2/3] Improve handlePostQueryHook readability --- .../opentelemetry-instrumentation-pg/src/instrumentation.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index 0906ad2c6f..4dbbe33315 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts @@ -154,6 +154,8 @@ export class PgInstrumentation extends InstrumentationBase { ); } + utils.handlePostQueryHook(pluginConfig, postQueryHookParams); + // Bind callback to parent span if (args.length > 0) { const parentSpan = trace.getSpan(context.active()); @@ -191,8 +193,6 @@ export class PgInstrumentation extends InstrumentationBase { } } - utils.handlePostQueryHook(pluginConfig, postQueryHookParams); - // Perform the original query const result: unknown = original.apply(this, args as never); From 27a405076c3bf2214a6d8d9f34227a73e8b8cde1 Mon Sep 17 00:00:00 2001 From: Alon Katz Date: Tue, 30 Nov 2021 16:47:43 +0200 Subject: [PATCH 3/3] Improve readability --- plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts index 7fb7daf451..1f93f286ad 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts @@ -595,6 +595,7 @@ describe('pg', () => { }); assert.strictEqual(called, true); }); + it('should call postQueryHook with query text and params if set', async () => { instrumentation.disable(); let called = false; @@ -627,6 +628,7 @@ describe('pg', () => { }); assert.strictEqual(called, true); }); + it('should call postQueryHook with query config if set', async () => { instrumentation.disable(); const name = 'fetch-text';