From f2faeda2993536547fecbbec63a4bc7ad83583e6 Mon Sep 17 00:00:00 2001
From: Joel Mukuthu <joelmukuthu@gmail.com>
Date: Mon, 3 Apr 2023 10:46:59 +0200
Subject: [PATCH] feat(instrumentation-pg): support custom SQL Commenter
 attributes

---
 .../src/utils.ts                              | 15 ++++--
 .../test/utils.test.ts                        | 54 +++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
index 2e4176817d..883086b451 100644
--- a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
+++ b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
@@ -47,6 +47,8 @@ function arrayStringifyHelper(arr: Array<unknown>): string {
   return '[' + arr.toString() + ']';
 }
 
+type SqlCommenterAttributes = Record<string, string>;
+
 /**
  * Helper function to get a low cardinality span name from whatever info we have
  * about the query.
@@ -284,7 +286,11 @@ function fixedEncodeURIComponent(str: string) {
   );
 }
 
-export function addSqlCommenterComment(span: Span, query: string): string {
+export function addSqlCommenterComment(
+  span: Span,
+  query: string,
+  customAttributes: SqlCommenterAttributes = {}
+): string {
   if (typeof query !== 'string' || query.length === 0) {
     return query;
   }
@@ -296,7 +302,7 @@ export function addSqlCommenterComment(span: Span, query: string): string {
   }
 
   const propagator = new W3CTraceContextPropagator();
-  const headers: { [key: string]: string } = {};
+  const headers: SqlCommenterAttributes = customAttributes;
   propagator.inject(
     trace.setSpan(ROOT_CONTEXT, span),
     headers,
@@ -311,9 +317,10 @@ export function addSqlCommenterComment(span: Span, query: string): string {
   }
 
   const commentString = sortedKeys
-    .map(key => {
+    .map((key) => {
+      const encodedKey = fixedEncodeURIComponent(key);
       const encodedValue = fixedEncodeURIComponent(headers[key]);
-      return `${key}='${encodedValue}'`;
+      return `${encodedKey}='${encodedValue}'`;
     })
     .join(',');
 
diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts
index 2e6c8ca0d7..cf657ff415 100644
--- a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts
@@ -289,5 +289,59 @@ describe('utils.ts', () => {
         "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01',tracestate='foo%3D%27bar%2Cbaz%3D%27qux%21%28%29%2A%27%2Chack%3D%27DROP%20TABLE'*/"
       );
     });
+
+    it('supports custom attributes', () => {
+      const spanContext: SpanContext = {
+        traceId: 'd4cda95b652f4a1592b449d5929fda1b',
+        spanId: '6e0c63257de34c92',
+        traceFlags: TraceFlags.SAMPLED,
+      };
+
+      const query = 'SELECT * from FOO;';
+      assert.strictEqual(
+        utils.addSqlCommenterComment(
+          trace.wrapSpanContext(spanContext),
+          query,
+          { controller: 'foo', action: 'bar' }
+        ),
+        "SELECT * from FOO; /*action='bar',controller='foo',traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/"
+      );
+    });
+
+    it('escapes special characters in keys', () => {
+      const spanContext: SpanContext = {
+        traceId: 'd4cda95b652f4a1592b449d5929fda1b',
+        spanId: '6e0c63257de34c92',
+        traceFlags: TraceFlags.SAMPLED,
+      };
+
+      const query = 'SELECT * from FOO;';
+      assert.strictEqual(
+        utils.addSqlCommenterComment(
+          trace.wrapSpanContext(spanContext),
+          query,
+          { "'DROP TABLE": 'foo' }
+        ),
+        "SELECT * from FOO; /*%27DROP%20TABLE='foo',traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/"
+      );
+    });
+
+    it('overwrites custom attributes with trace context keys', () => {
+      const spanContext: SpanContext = {
+        traceId: 'd4cda95b652f4a1592b449d5929fda1b',
+        spanId: '6e0c63257de34c92',
+        traceFlags: TraceFlags.SAMPLED,
+      };
+
+      const query = 'SELECT * from FOO;';
+      assert.strictEqual(
+        utils.addSqlCommenterComment(
+          trace.wrapSpanContext(spanContext),
+          query,
+          { traceparent: 'foo' }
+        ),
+        "SELECT * from FOO; /*traceparent='00-d4cda95b652f4a1592b449d5929fda1b-6e0c63257de34c92-01'*/"
+      );
+    });
   });
 });