feat(pg): support requestHook hook (#1307)

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

@@ -47,7 +47,8 @@ PostgreSQL instrumentation has few options available to choose from. You can set
 | Options | Type | Description |
 | ------- | ---- | ----------- |
 | [`enhancedDatabaseReporting`](./src/types.ts#L30) | `boolean` | If true, additional information about query parameters and results will be attached (as `attributes`) to spans representing database operations |
-| `responseHook` | `PgInstrumentationExecutionResponseHook` (function) | Function for adding custom attributes from db response |
+| `requestHook` | `PgInstrumentationExecutionRequestHook` (function) | Function for adding custom span attributes using information about the query being issued and the db to which it's directed |
+| `responseHook` | `PgInstrumentationExecutionResponseHook` (function) | Function for adding custom span attributes from db response |
 | `requireParentSpan` | `boolean` | If true, requires a parent span to create new spans (default false) |
 | `addSqlCommenterCommentToQueries` | `boolean` | If true, adds [sqlcommenter](https://github.com/open-telemetry/opentelemetry-sqlcommenter) specification compliant comment to queries with tracing context (default false). _NOTE: A comment will not be added to queries that already contain `--` or `/* ... */` in them, even if these are not actually part of comments_ |
 

plugins/node/opentelemetry-instrumentation-pg/package.json

@@ -68,6 +68,7 @@
     "pg": "8.7.1",
     "pg-pool": "3.4.1",
     "rimraf": "3.0.2",
+    "safe-stable-stringify": "^2.4.1",
     "sinon": "14.0.0",
     "test-all-versions": "5.0.1",
     "ts-mocha": "10.0.0",

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';
@@ -257,6 +258,50 @@ export class PgInstrumentation extends InstrumentationBase {
           }
         }
 
+        if (
+          typeof instrumentationConfig.requestHook === 'function' &&
+          queryConfig
+        ) {
+          safeExecuteInTheMiddle(
+            () => {
+              // pick keys to expose explicitly, so we're not leaking pg package
+              // internals that are subject to change
+              const { database, host, port, user } = this.connectionParameters;
+              const connection = { database, host, port, user };
+
+              instrumentationConfig.requestHook!(span, {
+                connection,
+                query: {
+                  text: queryConfig.text,
+                  // nb: if `client.query` is called with illegal arguments
+                  // (e.g., if `queryConfig.values` is passed explicitly, but a
+                  // non-array is given), then the type casts will be wrong. But
+                  // we leave it up to the queryHook to handle that, and we
+                  // catch and swallow any errors it throws. The other options
+                  // are all worse. E.g., we could leave `queryConfig.values`
+                  // and `queryConfig.name` as `unknown`, but then the hook body
+                  // would be forced to validate (or cast) them before using
+                  // them, which seems incredibly cumbersome given that these
+                  // casts will be correct 99.9% of the time -- and pg.query
+                  // will immediately throw during development in the other .1%
+                  // of cases. Alternatively, we could simply skip calling the
+                  // hook when `values` or `name` don't have the expected type,
+                  // but that would add unnecessary validation overhead to every
+                  // hook invocation and possibly be even more confusing/unexpected.
+                  values: queryConfig.values as unknown[],
+                  name: queryConfig.name as string | undefined,
+                },
+              });
+            },
+            err => {
+              if (err) {
+                plugin._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/internal-types.ts

@@ -19,17 +19,25 @@ import type * as pgPoolTypes from 'pg-pool';
 
 export type PostgresCallback = (err: Error, res: object) => unknown;
 
-// These are not included in @types/pg, so manually define them.
-// https://github.com/brianc/node-postgres/blob/fde5ec586e49258dfc4a2fcd861fcdecb4794fc3/lib/client.js#L25
-export interface PgClientConnectionParams {
+// NB: this type describes the shape of a parsed, normalized form of the
+// connection information that's stored inside each pg.Client instance. It's
+// _not_ the same as the ConnectionConfig type exported from `@types/pg`. That
+// type defines how data must be _passed in_ when creating a new `pg.Client`,
+// which doesn't necessarily match the normalized internal form. E.g., a user
+// can call `new Client({ connectionString: '...' }), but `connectionString`
+// will never show up in the type below, because only the extracted host, port,
+// etc. are recorded in this normalized config. The keys listed below are also
+// incomplete, which is fine because the type is internal and these keys are the
+// only ones our code is reading. See https://github.com/brianc/node-postgres/blob/fde5ec586e49258dfc4a2fcd861fcdecb4794fc3/lib/client.js#L25
+export interface PgParsedConnectionParams {
   database?: string;
   host?: string;
   port?: number;
   user?: string;
 }
 
 export interface PgClientExtended extends pgTypes.Client {
-  connectionParameters: PgClientConnectionParams;
+  connectionParameters: PgParsedConnectionParams;
 }
 
 export type PgPoolCallback = (

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

@@ -26,12 +26,39 @@ export interface PgInstrumentationExecutionResponseHook {
   (span: api.Span, responseInfo: PgResponseHookInformation): void;
 }
 
+export interface PgRequestHookInformation {
+  query: {
+    text: string;
+    name?: string;
+    values?: unknown[];
+  };
+  connection: {
+    database?: string;
+    host?: string;
+    port?: number;
+    user?: string;
+  };
+}
+
+export interface PgInstrumentationExecutionRequestHook {
+  (span: api.Span, queryInfo: PgRequestHookInformation): void;
+}
+
 export interface PgInstrumentationConfig extends InstrumentationConfig {
   /**
-   * If true, additional information about query parameters will be attached (as `attributes`) to spans representing
+   * If true, an attribute containing the query's parameters will be attached
+   * the spans generated to represent the query.
    */
   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
+   */
+  requestHook?: PgInstrumentationExecutionRequestHook;
+
   /**
    * Hook that allows adding custom span attributes based on the data
    * returned from "query" Pg actions.

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

@@ -36,10 +36,10 @@ import {
 import {
   PgClientExtended,
   PostgresCallback,
-  PgClientConnectionParams,
   PgErrorCallback,
   PgPoolCallback,
   PgPoolExtended,
+  PgParsedConnectionParams,
 } from './internal-types';
 import { PgInstrumentationConfig } from './types';
 import type * as pgTypes from 'pg';
@@ -97,15 +97,15 @@ function parseNormalizedOperationName(queryText: string) {
   return sqlCommand.endsWith(';') ? sqlCommand.slice(0, -1) : sqlCommand;
 }
 
-function getConnectionString(params: PgClientConnectionParams) {
+export function getConnectionString(params: PgParsedConnectionParams) {
   const host = params.host || 'localhost';
   const port = params.port || 5432;
   const database = params.database || '';
   return `postgresql://${host}:${port}/${database}`;
 }
 
 export function getSemanticAttributesFromConnection(
-  params: PgClientConnectionParams
+  params: PgParsedConnectionParams
 ) {
   return {
     [SemanticAttributes.DB_NAME]: params.database, // required

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

@@ -33,6 +33,7 @@ import {
 import * as assert from 'assert';
 import type * as pg from 'pg';
 import * as sinon from 'sinon';
+import stringify from 'safe-stable-stringify';
 import {
   PgInstrumentation,
   PgInstrumentationConfig,
@@ -513,6 +514,100 @@ describe('pg', () => {
       });
     });
 
+    describe('when specifying a requestHook configuration', () => {
+      const dataAttributeName = 'pg_data';
+      const query = 'SELECT 0::text';
+      const events: TimedEvent[] = [];
+
+      // these are the attributes that we'd expect would end up on the final
+      // span if there is no requestHook.
+      const attributes = {
+        ...DEFAULT_ATTRIBUTES,
+        [SemanticAttributes.DB_STATEMENT]: query,
+      };
+
+      // These are the attributes we expect on the span after the requestHook
+      // has run. We set up the hook to just add to the span a stringified
+      // version of the args it receives (which is an easy way to assert both
+      // that the proper args were passed and that the hook was called).
+      const attributesAfterHook = {
+        ...attributes,
+        [dataAttributeName]: stringify({
+          connection: {
+            database: CONFIG.database,
+            port: CONFIG.port,
+            host: CONFIG.host,
+            user: CONFIG.user,
+          },
+          query: { text: query },
+        }),
+      };
+
+      describe('AND valid requestHook', () => {
+        beforeEach(async () => {
+          create({
+            enhancedDatabaseReporting: true,
+            requestHook: (span, requestInfo) => {
+              span.setAttribute(dataAttributeName, stringify(requestInfo));
+            },
+          });
+        });
+
+        it('should attach request hook data to resulting spans for query with callback ', done => {
+          const span = tracer.startSpan('test span');
+          context.with(trace.setSpan(context.active(), span), () => {
+            const res = client.query(query, (err, res) => {
+              assert.strictEqual(err, null);
+              assert.ok(res);
+              runCallbackTest(span, attributesAfterHook, events);
+              done();
+            });
+            assert.strictEqual(res, undefined, 'No promise is returned');
+          });
+        });
+
+        it('should attach request hook data to resulting spans for query returning a Promise', async () => {
+          const span = tracer.startSpan('test span');
+          await context.with(
+            trace.setSpan(context.active(), span),
+            async () => {
+              const resPromise = await client.query({ text: query });
+              try {
+                assert.ok(resPromise);
+                runCallbackTest(span, attributesAfterHook, events);
+              } catch (e) {
+                assert.ok(false, e.message);
+              }
+            }
+          );
+        });
+      });
+
+      describe('AND invalid requestHook', () => {
+        beforeEach(async () => {
+          create({
+            enhancedDatabaseReporting: true,
+            requestHook: (_span, _requestInfo) => {
+              throw 'some kind of failure!';
+            },
+          });
+        });
+
+        it('should not do any harm when throwing an exception', done => {
+          const span = tracer.startSpan('test span');
+          context.with(trace.setSpan(context.active(), span), () => {
+            const res = client.query(query, (err, res) => {
+              assert.strictEqual(err, null);
+              assert.ok(res);
+              runCallbackTest(span, attributes, events);
+              done();
+            });
+            assert.strictEqual(res, undefined, 'No promise is returned');
+          });
+        });
+      });
+    });
+
     describe('when specifying a responseHook configuration', () => {
       const dataAttributeName = 'pg_data';
       const query = 'SELECT 0::text';