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

@@ -72,6 +72,7 @@
     "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,56 @@ 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;
+}
+
+export function addSqlCommenterComment(span: Span, query: string): string {
+  if (query.length === 0 || typeof query !== 'string') {
+    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 escapedValue = headers[key].replace(/([^\\])(')/g, (_, prefix) => {
+        return `${prefix}\\'`;
+      });
+      return `${key}='${escapedValue}'`;
+    })
+    .join(",");
+
+  return `${query} /*${commentString}*/`;
+}

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

@@ -44,6 +44,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;
@@ -99,6 +100,7 @@ describe('pg', () => {
     instrumentation.enable();
   }
 
+  let executedQueries: (pg.Query & { text?: string })[];
   let postgres: typeof pg;
   let client: pg.Client;
   let instrumentation: PgInstrumentation;
@@ -141,8 +143,21 @@ describe('pg', () => {
     context.setGlobalContextManager(contextManager);
     instrumentation.setTracerProvider(provider);
 
+    executedQueries = [];
     postgres = require('pg');
     client = new postgres.Client(CONFIG);
+
+    // We need to track the actual queries that get executed to assert
+    // that the sqlcommenter logic works
+    const privateClient: {
+      queryQueue: (pg.Query & { text?: string })[];
+    } = client as any;
+    const originalPush = privateClient.queryQueue.push;
+    privateClient.queryQueue.push = (...items) => {
+      executedQueries.push(...items);
+      return originalPush.apply(privateClient.queryQueue, items);
+    };
+
     await client.connect();
   });
 
@@ -161,6 +176,7 @@ describe('pg', () => {
   afterEach(() => {
     memoryExporter.reset();
     context.disable();
+    executedQueries.splice(0, executedQueries.length);
   });
 
   it('should return an instrumentation', () => {
@@ -649,6 +665,108 @@ 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
+          );
+          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
+            );
+            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
+          );
+          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
+            );
+            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,