fix(instrumentation-pg): preserve SQL template text property (#3259)

Author: aryamohanan

packages/instrumentation-pg/src/instrumentation.ts

@@ -333,19 +333,24 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
         // TODO: remove the `as ...` casts below when the TS version is upgraded.
         // Newer TS versions will use the result of firstArgIsQueryObjectWithText
         // to properly narrow arg0, but TS 4.3.5 does not.
-        const queryConfig = firstArgIsString
-          ? {
-              text: arg0 as string,
-              values: Array.isArray(args[1]) ? args[1] : undefined,
-            }
-          : firstArgIsQueryObjectWithText
-            ? {
-                ...(arg0 as any),
-                values:
-                  (arg0 as any).values ??
-                  (Array.isArray(args[1]) ? args[1] : undefined),
-              }
-            : undefined;
+        let queryConfig: any;
+
+        if (firstArgIsString) {
+          queryConfig = {
+            text: arg0 as string,
+            values: Array.isArray(args[1]) ? args[1] : undefined,
+          };
+        } else if (firstArgIsQueryObjectWithText) {
+          const q = arg0 as any;
+
+          if (q.values === undefined && Array.isArray(args[1])) {
+            q.values = args[1];
+          }
+
+          queryConfig = q;
+        } else {
+          queryConfig = undefined;
+        }
 
         const attributes: Attributes = {
           [ATTR_DB_SYSTEM]: DB_SYSTEM_VALUE_POSTGRESQL,

packages/instrumentation-pg/test/pg.test.ts

@@ -1039,6 +1039,101 @@ describe('pg', () => {
         done();
       });
     });
+
+    it('should generate traces for SQL template literal objects ', async () => {
+      const span = tracer.startSpan('test span');
+      await context.with(trace.setSpan(context.active(), span), async () => {
+        // Mock SQLStatement from sql-template-strings library
+        class SQLStatement {
+          strings: string[];
+          values: any[];
+
+          constructor(strings: string[], values: any[]) {
+            this.strings = strings;
+            this.values = values;
+          }
+
+          get text(): string {
+            return this.strings.reduce(
+              (prev: string, curr: string, i: number) => prev + '$' + i + curr
+            );
+          }
+        }
+
+        const name = 'pg_tables';
+        const sqlQuery = new SQLStatement(
+          ['SELECT * FROM information_schema.tables WHERE table_name = ', ''],
+          [name]
+        );
+        const resPromise = await client.query(sqlQuery);
+        assert.ok(resPromise);
+
+        const spans = memoryExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 1);
+        const pgSpan = spans[0];
+
+        assert.strictEqual(pgSpan.name, 'pg.query:SELECT otel_pg_database');
+        assert.strictEqual(
+          pgSpan.attributes[ATTR_DB_STATEMENT],
+          'SELECT * FROM information_schema.tables WHERE table_name = $1'
+        );
+
+        const executedQueries = getExecutedQueries();
+        assert.ok(executedQueries.length > 0);
+        assert.strictEqual(
+          executedQueries[executedQueries.length - 1].text,
+          'SELECT * FROM information_schema.tables WHERE table_name = $1'
+        );
+      });
+    });
+
+    it('should generate traces for SQL template literals with added comments', async () => {
+      instrumentation.setConfig({
+        addSqlCommenterCommentToQueries: true,
+      });
+
+      const span = tracer.startSpan('test span');
+
+      await context.with(trace.setSpan(context.active(), span), async () => {
+        // Mock SQLStatement from sql-template-strings library
+        class SQLStatement {
+          strings: string[];
+          values: any[];
+
+          constructor(strings: string[], values: any[]) {
+            this.strings = strings;
+            this.values = values;
+          }
+
+          get text(): string {
+            return this.strings.reduce(
+              (prev: string, curr: string, i: number) => prev + '$' + i + curr
+            );
+          }
+        }
+        const tableName = 'pg_tables';
+        const sqlQuery = new SQLStatement(
+          ['SELECT * FROM information_schema.tables WHERE table_name = ', ''],
+          [tableName]
+        );
+
+        const res = await client.query(sqlQuery);
+        assert.ok(res);
+
+        const spans = memoryExporter.getFinishedSpans();
+        assert.strictEqual(spans.length, 1);
+        const pgSpan = spans[0];
+
+        const commentedQuery = addSqlCommenterComment(
+          trace.wrapSpanContext(pgSpan.spanContext()),
+          sqlQuery.text
+        );
+        const executedQueries = getExecutedQueries();
+        assert.ok(executedQueries.length > 0);
+        const lastQuery = executedQueries[executedQueries.length - 1];
+        assert.strictEqual(lastQuery.text, commentedQuery);
+      });
+    });
   });
 
   describe('exception event recording', () => {