From 8f30771f8c420053838231a18e633e082a4b8ea5 Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad <aprasad@sentry.io>
Date: Wed, 9 Aug 2023 23:30:27 -0400
Subject: [PATCH] feat(tracing): Add db connection attributes for postgres
 spans

---
 .../suites/tracing/auto-instrument/pg/test.ts | 17 +++++++++
 .../src/node/integrations/postgres.ts         | 35 ++++++++++++++++---
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/packages/node-integration-tests/suites/tracing/auto-instrument/pg/test.ts b/packages/node-integration-tests/suites/tracing/auto-instrument/pg/test.ts
index 359e04d7d5f0..a9d4b56b0b3c 100644
--- a/packages/node-integration-tests/suites/tracing/auto-instrument/pg/test.ts
+++ b/packages/node-integration-tests/suites/tracing/auto-instrument/pg/test.ts
@@ -1,6 +1,11 @@
 import { assertSentryTransaction, TestEnv } from '../../../../utils';
 
 class PgClient {
+  database?: string = 'test';
+  user?: string = 'user';
+  host?: string = 'localhost';
+  port?: number = 5432;
+
   // https://node-postgres.com/api/client#clientquery
   public query(_text: unknown, values: unknown, callback?: () => void) {
     if (typeof callback === 'function') {
@@ -42,6 +47,10 @@ test('should auto-instrument `pg` package.', async () => {
         op: 'db',
         data: {
           'db.system': 'postgresql',
+          'db.user': 'user',
+          'db.name': 'test',
+          'server.address': 'localhost',
+          'server.port': 5432,
         },
       },
       {
@@ -49,6 +58,10 @@ test('should auto-instrument `pg` package.', async () => {
         op: 'db',
         data: {
           'db.system': 'postgresql',
+          'db.user': 'user',
+          'db.name': 'test',
+          'server.address': 'localhost',
+          'server.port': 5432,
         },
       },
       {
@@ -56,6 +69,10 @@ test('should auto-instrument `pg` package.', async () => {
         op: 'db',
         data: {
           'db.system': 'postgresql',
+          'db.user': 'user',
+          'db.name': 'test',
+          'server.address': 'localhost',
+          'server.port': 5432,
         },
       },
     ],
diff --git a/packages/tracing-internal/src/node/integrations/postgres.ts b/packages/tracing-internal/src/node/integrations/postgres.ts
index 85c021c4f24b..f226c9264938 100644
--- a/packages/tracing-internal/src/node/integrations/postgres.ts
+++ b/packages/tracing-internal/src/node/integrations/postgres.ts
@@ -11,6 +11,13 @@ interface PgClient {
   };
 }
 
+interface PgClientThis {
+  database?: string;
+  host?: string;
+  port?: number;
+  user?: string;
+}
+
 interface PgOptions {
   usePgNative?: boolean;
 }
@@ -74,15 +81,35 @@ export class Postgres implements LazyLoadedIntegration<PGModule> {
      * function (pg.Cursor) => pg.Cursor
      */
     fill(Client.prototype, 'query', function (orig: () => void | Promise<unknown>) {
-      return function (this: unknown, config: unknown, values: unknown, callback: unknown) {
+      return function (this: PgClientThis, config: unknown, values: unknown, callback: unknown) {
         const scope = getCurrentHub().getScope();
         const parentSpan = scope?.getSpan();
+
+        const data: Record<string, string | number> = {
+          'db.system': 'postgresql',
+        };
+
+        try {
+          if (this.database) {
+            data['db.name'] = this.database;
+          }
+          if (this.host) {
+            data['server.address'] = this.host;
+          }
+          if (this.port) {
+            data['server.port'] = this.port;
+          }
+          if (this.user) {
+            data['db.user'] = this.user;
+          }
+        } catch (e) {
+          // ignore
+        }
+
         const span = parentSpan?.startChild({
           description: typeof config === 'string' ? config : (config as { text: string }).text,
           op: 'db',
-          data: {
-            'db.system': 'postgresql',
-          },
+          data,
         });
 
         if (typeof callback === 'function') {