feat: add sqlcommenter comment with trace context to queries in pg instrumentation

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

@@ -48,6 +48,8 @@ PostgreSQL instrumentation has few options available to choose from. You can set
 | ------- | ---- | ----------- |
 | [`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 |
+| `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 |
 
 ## Useful links
 

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

@@ -60,18 +60,21 @@
     "@opentelemetry/sdk-trace-node": "^1.8.0",
     "@types/mocha": "7.0.2",
     "@types/node": "18.11.7",
+    "@types/sinon": "10.0.2",
     "cross-env": "7.0.3",
     "gts": "3.1.0",
     "mocha": "7.2.0",
     "nyc": "15.1.0",
     "pg": "8.7.1",
     "pg-pool": "3.4.1",
     "rimraf": "3.0.2",
+    "sinon": "14.0.0",
     "test-all-versions": "5.0.1",
     "ts-mocha": "10.0.0",
     "typescript": "4.3.5"
   },
   "dependencies": {
+    "@opentelemetry/core": "^1.8.0",
     "@opentelemetry/instrumentation": "^0.34.0",
     "@opentelemetry/semantic-conventions": "^1.0.0",
     "@types/pg": "8.6.1",

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

@@ -178,6 +178,7 @@ export class PgInstrumentation extends InstrumentationBase {
         // 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(
@@ -195,14 +196,29 @@ export class PgInstrumentation extends InstrumentationBase {
               query
             );
           }
+
+          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,

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

@@ -46,4 +46,10 @@ export interface PgInstrumentationConfig extends InstrumentationConfig {
    * @default false
    */
   requireParentSpan?: boolean;
+
+  /**
+   * If true, queries are modified to also include a comment with
+   * the tracing context, following the {@link https://github.com/open-telemetry/opentelemetry-sqlcommenter sqlcommenter} format
+   */
+  addSqlCommenterCommentToQueries?: boolean;
 }

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

@@ -24,7 +24,10 @@ import {
   diag,
   INVALID_SPAN_CONTEXT,
   Attributes,
+  defaultTextMapSetter,
+  ROOT_CONTEXT,
 } from '@opentelemetry/api';
+import { W3CTraceContextPropagator } from '@opentelemetry/core';
 import { AttributeNames } from './enums/AttributeNames';
 import {
   SemanticAttributes,
@@ -289,3 +292,65 @@ export function patchClientConnectCallback(
     cb.call(this, err);
   };
 }
+
+function hasValidSqlComment(query: string): boolean {
+  const indexOpeningDashDashComment = query.indexOf('--');
+  if (indexOpeningDashDashComment >= 0) {
+    return true;
+  }
+
+  const indexOpeningSlashComment = query.indexOf('/*');
+  if (indexOpeningSlashComment < 0) {
+    return false;
+  }
+
+  const indexClosingSlashComment = query.indexOf('*/');
+  return indexOpeningDashDashComment < indexClosingSlashComment;
+}
+
+// sqlcommenter specification expects us to URL encode all reserved
+// characters, but encodeURIComponent does not handle ! ' ( ) *,
+// which means we need special handling for this
+// https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/encodeURIComponent
+function fixedEncodeURIComponent(str: string) {
+  return encodeURIComponent(str).replace(
+    /[!'()*]/g,
+    c => `%${c.charCodeAt(0).toString(16).toUpperCase()}`
+  );
+}
+
+export function addSqlCommenterComment(span: Span, query: string): string {
+  if (typeof query !== 'string' || query.length === 0) {
+    return query;
+  }
+
+  // As per sqlcommenter spec we shall not add a comment if there already is a comment
+  // in the query
+  if (hasValidSqlComment(query)) {
+    return query;
+  }
+
+  const propagator = new W3CTraceContextPropagator();
+  const headers: { [key: string]: string } = {};
+  propagator.inject(
+    trace.setSpan(ROOT_CONTEXT, span),
+    headers,
+    defaultTextMapSetter
+  );
+
+  // sqlcommenter spec requires keys in the comment to be sorted lexicographically
+  const sortedKeys = Object.keys(headers).sort();
+
+  if (sortedKeys.length === 0) {
+    return query;
+  }
+
+  const commentString = sortedKeys
+    .map(key => {
+      const encodedValue = fixedEncodeURIComponent(headers[key]);
+      return `${key}='${encodedValue}'`;
+    })
+    .join(',');
+
+  return `${query} /*${commentString}*/`;
+}

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

@@ -32,6 +32,7 @@ import {
 } from '@opentelemetry/sdk-trace-base';
 import * as assert from 'assert';
 import type * as pg from 'pg';
+import * as Sinon from 'sinon';
 import {
   PgInstrumentation,
   PgInstrumentationConfig,
@@ -44,6 +45,7 @@ import {
   DbSystemValues,
 } from '@opentelemetry/semantic-conventions';
 import { isSupported } from './utils';
+import { addSqlCommenterComment } from '../src/utils';
 
 const pgVersion = require('pg/package.json').version;
 const nodeVersion = process.versions.node;
@@ -110,6 +112,12 @@ describe('pg', () => {
   const testPostgresLocally = process.env.RUN_POSTGRES_TESTS_LOCAL; // For local: spins up local postgres db via docker
   const shouldTest = testPostgres || testPostgresLocally; // Skips these tests if false (default)
 
+  function getExecutedQueries() {
+    return (client as any).queryQueue.push.args.flat() as (pg.Query & {
+      text?: string;
+    })[];
+  }
+
   before(async function () {
     const skipForUnsupported =
       process.env.IN_TAV && !isSupported(nodeVersion, pgVersion);
@@ -156,11 +164,16 @@ describe('pg', () => {
   beforeEach(() => {
     contextManager = new AsyncHooksContextManager().enable();
     context.setGlobalContextManager(contextManager);
+
+    // Add a spy on the underlying client's internal query queue so that
+    // we could assert on what the final queries are that are executed
+    Sinon.spy((client as any).queryQueue, 'push');
   });
 
   afterEach(() => {
     memoryExporter.reset();
     context.disable();
+    Sinon.restore();
   });
 
   it('should return an instrumentation', () => {
@@ -649,6 +662,116 @@ describe('pg', () => {
       });
     });
 
+    it('should not add sqlcommenter comment when flag is not specified', async () => {
+      const span = tracer.startSpan('test span');
+      await context.with(trace.setSpan(context.active(), span), async () => {
+        try {
+          const query = 'SELECT NOW()';
+          const resPromise = await client.query(query);
+          assert.ok(resPromise);
+
+          const [span] = memoryExporter.getFinishedSpans();
+          assert.ok(span);
+
+          const commentedQuery = addSqlCommenterComment(
+            trace.wrapSpanContext(span.spanContext()),
+            query
+          );
+
+          const executedQueries = getExecutedQueries();
+          assert.equal(executedQueries.length, 1);
+          assert.equal(executedQueries[0].text, query);
+          assert.notEqual(query, commentedQuery);
+        } catch (e) {
+          assert.ok(false, e.message);
+        }
+      });
+    });
+
+    it('should not add sqlcommenter comment with client.query({text, callback}) when flag is not specified', done => {
+      const span = tracer.startSpan('test span');
+      context.with(trace.setSpan(context.active(), span), () => {
+        const query = 'SELECT NOW()';
+        client.query({
+          text: query,
+          callback: (err: Error, res: pg.QueryResult) => {
+            assert.strictEqual(err, null);
+            assert.ok(res);
+
+            const [span] = memoryExporter.getFinishedSpans();
+            const commentedQuery = addSqlCommenterComment(
+              trace.wrapSpanContext(span.spanContext()),
+              query
+            );
+
+            const executedQueries = getExecutedQueries();
+            assert.equal(executedQueries.length, 1);
+            assert.equal(executedQueries[0].text, query);
+            assert.notEqual(query, commentedQuery);
+            done();
+          },
+        } as pg.QueryConfig);
+      });
+    });
+
+    it('should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified', async () => {
+      instrumentation.setConfig({
+        addSqlCommenterCommentToQueries: true,
+      });
+
+      const span = tracer.startSpan('test span');
+      await context.with(trace.setSpan(context.active(), span), async () => {
+        try {
+          const query = 'SELECT NOW()';
+          const resPromise = await client.query(query);
+          assert.ok(resPromise);
+
+          const [span] = memoryExporter.getFinishedSpans();
+          const commentedQuery = addSqlCommenterComment(
+            trace.wrapSpanContext(span.spanContext()),
+            query
+          );
+
+          const executedQueries = getExecutedQueries();
+          assert.equal(executedQueries.length, 1);
+          assert.equal(executedQueries[0].text, commentedQuery);
+          assert.notEqual(query, commentedQuery);
+        } catch (e) {
+          assert.ok(false, e.message);
+        }
+      });
+    });
+
+    it('should add sqlcommenter comment when addSqlCommenterCommentToQueries=true is specified with client.query({text, callback})', done => {
+      instrumentation.setConfig({
+        addSqlCommenterCommentToQueries: true,
+      });
+
+      const span = tracer.startSpan('test span');
+      context.with(trace.setSpan(context.active(), span), () => {
+        const query = 'SELECT NOW()';
+        client.query({
+          text: query,
+          callback: (err: Error, res: pg.QueryResult) => {
+            assert.strictEqual(err, null);
+            assert.ok(res);
+
+            const [span] = memoryExporter.getFinishedSpans();
+            const commentedQuery = addSqlCommenterComment(
+              trace.wrapSpanContext(span.spanContext()),
+              query
+            );
+
+            const executedQueries = getExecutedQueries();
+            assert.equal(executedQueries.length, 1);
+            assert.equal(executedQueries[0].text, commentedQuery);
+            assert.notEqual(query, commentedQuery);
+            done();
+          },
+        } as pg.QueryConfig);
+      });
+    });
+
     it('should not generate traces for client.query() when requireParentSpan=true is specified', done => {
       instrumentation.setConfig({
         requireParentSpan: true,