Skip to content

Commit

Permalink
refactor(pg): more defensive, DRYer implementation
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@ethanresnick
ethanresnick committed Nov 26, 2022
1 parent 4cafe6b commit bb2a586
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 191 deletions.
144 changes: 79 additions & 65 deletions plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@ import {
PgClientConnect,
PgClientExtended,
PgErrorCallback,
NormalizedQueryConfig,
PostgresCallback,
PgPoolExtended,
PgPoolCallback,
Expand Down Expand Up @@ -173,101 +172,107 @@ 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(
context.active(),
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) {
Expand Down Expand Up @@ -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 };
6 changes: 0 additions & 6 deletions plugins/node/opentelemetry-instrumentation-pg/src/internal-types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
98 changes: 21 additions & 77 deletions plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ import {
} from '@opentelemetry/semantic-conventions';
import {
PgClientExtended,
NormalizedQueryConfig,
PostgresCallback,
PgClientConnectionParams,
PgErrorCallback,
Expand Down Expand Up @@ -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,
Expand Down
Loading

0 comments on commit bb2a586

Please sign in to comment.