feat(instrumentation-pg): add post-query custom attributes hook for spans

plugins/node/opentelemetry-instrumentation-pg/README.md

@@ -50,6 +50,7 @@ PostgreSQL instrumentation has few options available to choose from. You can set
 | Options | Type | Description |
 | ------- | ---- | ----------- |
 | [`enhancedDatabaseReporting`](https:/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 |
 
 ## Useful links

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,
@@ -148,13 +154,15 @@ export class PgInstrumentation extends InstrumentationBase {
  );
  }
 
+ utils.handlePostQueryHook(pluginConfig, postQueryHookParams);
+
  // Bind callback to parent span
  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() as PgInstrumentationConfig,
+ pluginConfig,
  span,
  args[args.length - 1] as PostgresCallback
  );
@@ -170,7 +178,7 @@ export class PgInstrumentation extends InstrumentationBase {
  ) {
  // Patch ConfigQuery callback
  let callback = utils.patchCallback(
- plugin.getConfig() as PgInstrumentationConfig,
+ pluginConfig,
  span,
  (args[0] as NormalizedQueryConfig).callback!
  );
@@ -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);
  });

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.

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,

plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts

@@ -563,5 +563,113 @@ 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);
+ });
  });
 });