feat: support query hook

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

@@ -17,6 +17,7 @@ import {
   isWrapped,
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
+  safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 
 import { context, diag, trace, Span, SpanStatusCode } from '@opentelemetry/api';
@@ -240,6 +241,33 @@ export class PgInstrumentation extends InstrumentationBase {
           }
         }
 
+        if (
+          typeof instrumentationConfig.queryHook === 'function' &&
+          queryConfig
+        ) {
+          safeExecuteInTheMiddle(
+            () => {
+              instrumentationConfig.queryHook!(span, {
+                data: {
+                  text: queryConfig.text,
+                  // nb: the casts below are not type safe; they'll be wrong in
+                  // the presence of illegal arguments. But we leave it up to
+                  // the queryHook to handle that, and we catch and swallow any
+                  // errors it throws.
+                  values: queryConfig.values as unknown[],
+                  name: queryConfig.name as string | undefined,
+                },
+              });
+            },
+            err => {
+              if (err) {
+                diag.error('Error running query hook', err);
+              }
+            },
+            true
+          );
+        }
+
         let result: unknown;
         try {
           result = original.apply(this, args as never);

plugins/node/opentelemetry-instrumentation-pg/src/types.ts

@@ -26,12 +26,28 @@ export interface PgInstrumentationExecutionResponseHook {
   (span: api.Span, responseInfo: PgResponseHookInformation): void;
 }
 
+export interface PgQueryHookInformation {
+  data: pgTypes.QueryConfig;
+}
+
+export interface PgInstrumentationExecutionQueryHook {
+  (span: api.Span, queryInfo: PgQueryHookInformation): void;
+}
+
 export interface PgInstrumentationConfig extends InstrumentationConfig {
   /**
    * If true, additional information about query parameters will be attached (as `attributes`) to spans representing
    */
   enhancedDatabaseReporting?: boolean;
 
+  /**
+   * Hook that allows adding custom span attributes or updating the
+   * span's name based on the data about the query to execute.
+   *
+   * @default undefined
+   */
+  queryHook?: PgInstrumentationExecutionQueryHook;
+
   /**
    * Hook that allows adding custom span attributes based on the data
    * returned from "query" Pg actions.