From bb2a5865337f2013bac1926e428ecae3ef311436 Mon Sep 17 00:00:00 2001 From: Ethan Resnick Date: Fri, 25 Nov 2022 19:37:01 -0800 Subject: [PATCH] refactor(pg): more defensive, DRYer implementation The prior code had three different utility functions for constructing the span to use to trace the query: handleTextQuery, handleConfigQuery, and handleParameterizedQuery. These functions contained a lot of duplicated logic, and that duplication was gonna make it harder to introduce new features (see next commit) and was already leading to bugs/inconsistencies. For example, `handleConfigQuery` would verify that the parameter `values` provided with the query were provided as an array before attaching them to the span, whereas `handleParameterizedQuery` would do so without this validation. This commit instead normalizes the arguments to `client.query` first, so that only function is needed to create the span. In the process of normalizing the arguments, the new code also performs a bit more validation on them, to make sure that the instrumentation will not throw at runtime. For example, the prior code would've thrown on `client.query(null)`, as it was checking `typeof args[0] === 'object'` without excluding null. The prior code also assumed, if an object was passed, that the `text` and `name` keys held a string; this is now verified before attaching that data to the span. Finally, the prior code had two different code paths for invoking the original `client.query` method; one that went through `handleInvalidQuery` and one directly in `_getClientQueryPatch`. The former wrapped the call in a try-catch, while the latter did not. Now, there's only one call to the original client.query, and it's always in a try-catch, which ensures that other cases of invalid args passed to client.query (which might lead it to throw synchronously) won't be an issue. --- .../src/instrumentation.ts | 144 ++++++++++-------- .../src/internal-types.ts | 6 - .../src/utils.ts | 98 +++--------- .../test/utils.test.ts | 45 +----- 4 files changed, 102 insertions(+), 191 deletions(-) diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index f309555951..ae53a6531d 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts @@ -26,7 +26,6 @@ import { PgClientConnect, PgClientExtended, PgErrorCallback, - NormalizedQueryConfig, PostgresCallback, PgPoolExtended, PgPoolCallback, @@ -173,72 +172,66 @@ export class PgInstrumentation extends InstrumentationBase { `Patching ${PgInstrumentation.COMPONENT}.Client.prototype.query` ); return function query(this: PgClientExtended, ...args: unknown[]) { - let span: Span; - - // 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( - this, - plugin.tracer, - plugin.getConfig(), - query, - params - ); - } else { - span = utils.handleTextQuery.call( - this, - plugin.tracer, - plugin.getConfig(), - query - ); - } + // client.query(text, cb?), client.query(text, values, cb?), and + // client.query(configObj, cb?) are all valid signatures. We construct + // a queryConfig obj from all (valid) signatures to build the span in a + // unified way. We verify that we at least have query text, and code + // defensively when dealing with `queryConfig` after that (to handle all + // the other invalid cases, like a non-array for values being provided). + // The type casts here reflect only what we've actually validated. + const arg0 = args[0]; + const firstArgIsString = typeof arg0 === 'string'; + const firstArgIsQueryObjectWithText = isObjectWithTextString(arg0); + + // TODO: remove the `as ...` casts 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. + const queryConfig = firstArgIsString + ? { text: arg0 as string, values: args[1] } + : firstArgIsQueryObjectWithText + ? (arg0 as ObjectWithText) + : undefined; + + const instrumentationConfig = plugin.getConfig(); + + const span = utils.handleConfigQuery.call( + this, + plugin.tracer, + instrumentationConfig, + queryConfig + ); - 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, - plugin.tracer, - plugin.getConfig(), - original, - ...args - ); + // 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. + // Newer TS versions will use the result of firstArgIsQueryObjectWithText + // to properly narrow arg0, but TS 4.3.5 does not. + if (instrumentationConfig.addSqlCommenterCommentToQueries) { + args[0] = firstArgIsString + ? utils.addSqlCommenterComment(span, arg0 as string) + : firstArgIsQueryObjectWithText + ? { + ...(arg0 as ObjectWithText), + text: utils.addSqlCommenterComment( + span, + (arg0 as ObjectWithText).text + ), + } + : args[0]; } - // Bind callback to parent span + // Bind callback (if any) to parent span (if any) if (args.length > 0) { const parentSpan = trace.getSpan(context.active()); if (typeof args[args.length - 1] === 'function') { // Patch ParameterQuery callback args[args.length - 1] = utils.patchCallback( - plugin.getConfig(), + instrumentationConfig, span, - args[args.length - 1] as PostgresCallback + args[args.length - 1] as PostgresCallback // nb: not type safe. ); + // If a parent span exists, bind the callback if (parentSpan) { args[args.length - 1] = context.bind( @@ -246,28 +239,40 @@ export class PgInstrumentation extends InstrumentationBase { args[args.length - 1] ); } - } else if ( - typeof (args[0] as NormalizedQueryConfig).callback === 'function' - ) { + } else if (typeof queryConfig?.callback === 'function') { // Patch ConfigQuery callback let callback = utils.patchCallback( plugin.getConfig(), span, - (args[0] as NormalizedQueryConfig).callback! + queryConfig.callback as PostgresCallback // nb: not type safe. ); + // If a parent span existed, bind the callback if (parentSpan) { callback = context.bind(context.active(), callback); } - // Copy the callback instead of writing to args.callback so that we don't modify user's - // original callback reference + // Copy the callback instead of writing to args.callback so that we + // don't modify user's original callback reference args[0] = { ...(args[0] as object), callback }; } } - // Perform the original query - const result: unknown = original.apply(this, args as any); + let result: unknown; + try { + result = original.apply(this, args as never); + } catch (e) { + // span.recordException(e); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: + typeof e === 'object' && e !== null + ? String((e as { message?: unknown }).message) + : undefined, + }); + span.end(); + throw e; + } // Bind promise to parent span and end the span if (result instanceof Promise) { @@ -369,3 +374,12 @@ function handleConnectResult(span: Span, connectResult: unknown) { }) ); } + +function isObjectWithTextString(it: unknown): it is ObjectWithText { + return ( + typeof it === 'object' && + typeof (it as null | { text?: unknown })?.text === 'string' + ); +} + +type ObjectWithText = { text: string; [k: string]: unknown }; diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/internal-types.ts b/plugins/node/opentelemetry-instrumentation-pg/src/internal-types.ts index dc4e54fe7a..60b2377bc9 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/internal-types.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/internal-types.ts @@ -32,12 +32,6 @@ export interface PgClientExtended extends pgTypes.Client { connectionParameters: PgClientConnectionParams; } -// Interface name based on original driver implementation -// https://github.com/brianc/node-postgres/blob/2ef55503738eb2cbb6326744381a92c0bc0439ab/packages/pg/lib/utils.js#L157 -export interface NormalizedQueryConfig extends pgTypes.QueryConfig { - callback?: PostgresCallback; -} - export type PgPoolCallback = ( err: Error, client: any, diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts index 048c0d7ad0..140e841e02 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts @@ -35,7 +35,6 @@ import { } from '@opentelemetry/semantic-conventions'; import { PgClientExtended, - NormalizedQueryConfig, PostgresCallback, PgClientConnectionParams, PgErrorCallback, @@ -103,109 +102,54 @@ function startQuerySpan( }); } -// Queries where args[0] is a QueryConfig +// Create a span from our normalized queryConfig object, +// or return a basic span if no queryConfig was given/could be created. export function handleConfigQuery( this: PgClientExtended, tracer: Tracer, instrumentationConfig: PgInstrumentationConfig, - queryConfig: NormalizedQueryConfig + queryConfig?: { text: string; values?: unknown; name?: unknown } ) { + if (!queryConfig) { + return startQuerySpan( + this, + tracer, + instrumentationConfig, + PgInstrumentation.BASE_SPAN_NAME + ); + } + // Set child span name - const queryCommand = getCommandFromText(queryConfig.name || queryConfig.text); - const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand; + const queryCommand = getCommandFromText( + (typeof queryConfig.name === 'string' ? queryConfig.name : undefined) || + queryConfig.text + ); + const name = `${PgInstrumentation.BASE_SPAN_NAME}:${queryCommand}`; const span = startQuerySpan(this, tracer, instrumentationConfig, name); // Set attributes if (queryConfig.text) { span.setAttribute(SemanticAttributes.DB_STATEMENT, queryConfig.text); } + if ( instrumentationConfig.enhancedDatabaseReporting && - queryConfig.values instanceof Array + Array.isArray(queryConfig.values) ) { span.setAttribute( AttributeNames.PG_VALUES, arrayStringifyHelper(queryConfig.values) ); } + // Set plan name attribute, if present - if (queryConfig.name) { + if (typeof queryConfig.name === 'string') { span.setAttribute(AttributeNames.PG_PLAN, queryConfig.name); } return span; } -// Queries where args[1] is a 'values' array -export function handleParameterizedQuery( - this: PgClientExtended, - tracer: Tracer, - instrumentationConfig: PgInstrumentationConfig, - query: string, - values: unknown[] -) { - // Set child span name - const queryCommand = getCommandFromText(query); - const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand; - const span = startQuerySpan(this, tracer, instrumentationConfig, name); - - // Set attributes - span.setAttribute(SemanticAttributes.DB_STATEMENT, query); - if (instrumentationConfig.enhancedDatabaseReporting) { - span.setAttribute(AttributeNames.PG_VALUES, arrayStringifyHelper(values)); - } - - return span; -} - -// Queries where args[0] is a text query and 'values' was not specified -export function handleTextQuery( - this: PgClientExtended, - tracer: Tracer, - instrumentationConfig: PgInstrumentationConfig, - query: string -) { - // Set child span name - const queryCommand = getCommandFromText(query); - const name = PgInstrumentation.BASE_SPAN_NAME + ':' + queryCommand; - const span = startQuerySpan(this, tracer, instrumentationConfig, name); - - // Set attributes - span.setAttribute(SemanticAttributes.DB_STATEMENT, query); - - return span; -} - -/** - * Invalid query handler. We should never enter this function unless invalid args were passed to the driver. - * Create and immediately end a new span - */ -export function handleInvalidQuery( - this: PgClientExtended, - tracer: Tracer, - instrumentationConfig: PgInstrumentationConfig, - originalQuery: typeof pgTypes.Client.prototype.query, - ...args: unknown[] -) { - let result; - const span = startQuerySpan( - this, - tracer, - instrumentationConfig, - PgInstrumentation.BASE_SPAN_NAME - ); - try { - result = originalQuery.apply(this, args as never); - } catch (e) { - // span.recordException(e); - span.setStatus({ code: SpanStatusCode.ERROR, message: e.message }); - throw e; - } finally { - span.end(); - } - return result; -} - export function handleExecutionResult( config: PgInstrumentationConfig, span: Span, diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts index 4b43e4da95..1ab6dd0c4d 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts @@ -33,7 +33,7 @@ import * as assert from 'assert'; import * as pg from 'pg'; import { PgInstrumentationConfig } from '../src'; import { AttributeNames } from '../src/enums/AttributeNames'; -import { PgClientExtended, NormalizedQueryConfig } from '../src/internal-types'; +import { PgClientExtended } from '../src/internal-types'; import * as utils from '../src/utils'; const memoryExporter = new InMemorySpanExporter(); @@ -131,7 +131,7 @@ describe('utils.ts', () => { }); describe('.handleConfigQuery()', () => { - const queryConfig: NormalizedQueryConfig = { + const queryConfig = { text: 'SELECT $1::text', values: ['0'], }; @@ -171,47 +171,6 @@ describe('utils.ts', () => { }); }); - describe('.handleParameterizedQuery()', () => { - const query = 'SELECT $1::text'; - const values = ['0']; - - it('does not track pg.values by default', async () => { - const querySpan = utils.handleParameterizedQuery.call( - client, - tracer, - instrumentationConfig, - query, - values - ); - querySpan.end(); - - const readableSpan = getLatestSpan(); - - const pgValues = readableSpan.attributes[AttributeNames.PG_VALUES]; - assert.strictEqual(pgValues, undefined); - }); - - it('tracks pg.values if enabled explicitly', async () => { - const extPluginConfig: PgInstrumentationConfig & InstrumentationConfig = { - ...instrumentationConfig, - enhancedDatabaseReporting: true, - }; - const querySpan = utils.handleParameterizedQuery.call( - client, - tracer, - extPluginConfig, - query, - values - ); - querySpan.end(); - - const readableSpan = getLatestSpan(); - - const pgValues = readableSpan.attributes[AttributeNames.PG_VALUES]; - assert.strictEqual(pgValues, '[0]'); - }); - }); - describe('addSqlCommenterComment', () => { it('adds comment to a simple query', () => { const spanContext: SpanContext = {