add support for changing the pg service name based on connection para…

…ms (#1110)
DataDog · Oct 16, 2020 · 5a8b92b · 5a8b92b
1 parent aa9a458
commit 5a8b92b
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 4 deletions.
diff --git a/docs/test.ts b/docs/test.ts
@@ -204,6 +204,7 @@ tracer.use('net');
 tracer.use('paperplane');
 tracer.use('paperplane', httpServerOptions);
 tracer.use('pg');
+tracer.use('pg', { service: params => `${params.host}-${params.database}` });
 tracer.use('pino');
 tracer.use('promise-js');
 tracer.use('promise');

diff --git a/index.d.ts b/index.d.ts
@@ -524,7 +524,7 @@ declare namespace plugins {
     /**
      * The service name to be used for this plugin.
      */
-    service?: string;
+    service?: string | any;
 
     /** Whether to enable the plugin.
      * @default true
@@ -1027,7 +1027,12 @@ declare namespace plugins {
    * This plugin automatically instruments the
    * [pg](https://node-postgres.com/) module.
    */
-  interface pg extends Instrumentation {}
+  interface pg extends Instrumentation {
+    /**
+     * The service name to be used for this plugin. If a function is used, it will be passed the connection parameters and its return value will be used as the service name.
+     */
+    service?: string | ((params: any) => string);
+  }
 
   /**
    * This plugin patches the [pino](http://getpino.io)

diff --git a/packages/datadog-plugin-pg/src/index.js b/packages/datadog-plugin-pg/src/index.js
@@ -10,11 +10,13 @@ function createWrapQuery (tracer, config) {
     return function queryWithTrace () {
       const scope = tracer.scope()
       const childOf = scope.active()
+      const params = this.connectionParameters
+      const service = getServiceName(tracer, config, params)
       const span = tracer.startSpan(OPERATION_NAME, {
         childOf,
         tags: {
           [Tags.SPAN_KIND]: Tags.SPAN_KIND_RPC_CLIENT,
-          'service.name': config.service || `${tracer._service}-postgres`,
+          'service.name': service,
           'span.type': 'sql',
           'db.type': 'postgres'
         }
@@ -33,7 +35,6 @@ function createWrapQuery (tracer, config) {
 
       const originalCallback = pgQuery.callback
       const statement = pgQuery.text
-      const params = this.connectionParameters
 
       span.setTag('resource.name', statement)
 
@@ -69,6 +70,16 @@ function createWrapQuery (tracer, config) {
   }
 }
 
+function getServiceName (tracer, config, params) {
+  if (typeof config.service === 'function') {
+    return config.service(params)
+  } else if (config.service) {
+    return config.service
+  } else {
+    return `${tracer._service}-postgres`
+  }
+}
+
 module.exports = [
   {
     name: 'pg',

diff --git a/packages/datadog-plugin-pg/test/index.spec.js b/packages/datadog-plugin-pg/test/index.spec.js
@@ -196,6 +196,48 @@ describe('Plugin', () => {
           })
         })
       })
+
+      describe('with a service name callback', () => {
+        before(() => {
+          return agent.load('pg', { service: params => `${params.host}-${params.database}` })
+        })
+
+        after(() => {
+          return agent.close()
+        })
+
+        beforeEach(done => {
+          pg = require(`../../../versions/pg@${version}`).get()
+
+          client = new pg.Client({
+            user: 'postgres',
+            password: 'postgres',
+            database: 'postgres'
+          })
+
+          client.connect(err => done(err))
+        })
+
+        it('should be configured with the correct service', done => {
+          agent.use(traces => {
+            try {
+              expect(traces[0][0]).to.have.property('service', 'localhost-postgres')
+
+              done()
+            } catch (e) {
+              done(e)
+            }
+          })
+
+          client.query('SELECT $1::text as message', ['Hello world!'], (err, result) => {
+            if (err) throw err
+
+            client.end((err) => {
+              if (err) throw err
+            })
+          })
+        })
+      })
     })
   })
 })