tracing: add support for pg cursors / streams (#5680)

* add support for pg cursors / streams

Author: wconti27

packages/datadog-instrumentations/src/pg.js

@@ -14,6 +14,8 @@ const errorCh = channel('apm:pg:query:error')
 const startPoolQueryCh = channel('datadog:pg:pool:query:start')
 const finishPoolQueryCh = channel('datadog:pg:pool:query:finish')
 
+const { errorMonitor } = require('node:events')
+
 addHook({ name: 'pg', versions: ['>=8.0.3'] }, pg => {
   shimmer.wrap(pg.Client.prototype, 'query', query => wrapQuery(query))
   shimmer.wrap(pg.Pool.prototype, 'query', query => wrapPoolQuery(query))
@@ -39,13 +41,15 @@ function wrapQuery (query) {
       ? arguments[0]
       : { text: arguments[0] }
 
-    const textProp = Object.getOwnPropertyDescriptor(pgQuery, 'text')
+    const textPropObj = pgQuery.cursor ?? pgQuery
+    const textProp = Object.getOwnPropertyDescriptor(textPropObj, 'text')
+    const stream = typeof textPropObj.read === 'function'
 
     // Only alter `text` property if safe to do so.
     if (!textProp || textProp.configurable) {
-      const originalText = pgQuery.text
+      const originalText = textPropObj.text
 
-      Object.defineProperty(pgQuery, 'text', {
+      Object.defineProperty(textPropObj, 'text', {
         get () {
           return this?.__ddInjectableQuery || originalText
         }
@@ -57,9 +61,10 @@ function wrapQuery (query) {
 
       startCh.publish({
         params: this.connectionParameters,
-        query: pgQuery,
+        query: textPropObj,
         processId,
-        abortController
+        abortController,
+        stream
       })
 
       const finish = asyncResource.bind(function (error, res) {
@@ -130,6 +135,15 @@ function wrapQuery (query) {
         newQuery.then((res) => finish(null, res), finish)
       }
 
+      if (stream) {
+        newQuery.on('end', () => {
+          finish(null, [])
+        })
+        newQuery.on(errorMonitor, (err) => {
+          finish(err)
+        })
+      }
+
       try {
         return retval
       } catch (err) {

packages/datadog-plugin-pg/src/index.js

@@ -8,7 +8,7 @@ class PGPlugin extends DatabasePlugin {
   static get operation () { return 'query' }
   static get system () { return 'postgres' }
 
-  start ({ params = {}, query, processId }) {
+  start ({ params = {}, query, processId, stream }) {
     const service = this.serviceName({ pluginConfig: this.config, params })
     const originalStatement = this.maybeTruncate(query.text)
 
@@ -27,6 +27,10 @@ class PGPlugin extends DatabasePlugin {
       }
     })
 
+    if (stream) {
+      span.setTag('db.stream', 1)
+    }
+
     query.__ddInjectableQuery = this.injectDbmQuery(span, query.text, service, !!query.name)
   }
 }

packages/datadog-plugin-pg/test/index.spec.js

@@ -216,6 +216,73 @@ describe('Plugin', () => {
               .catch(done),
             rawExpectedSchema.outbound
           )
+
+          if (implementation !== 'pg.native') {
+            // pg-cursor is not supported on pg.native, pg-query-stream uses pg-cursor so it is also unsupported
+            describe('streaming capabilities', () => {
+              withVersions('pg', 'pg-cursor', pgCursorVersion => {
+                let Cursor
+
+                beforeEach(() => {
+                  Cursor = require(`../../../versions/pg-cursor@${pgCursorVersion}`).get()
+                })
+
+                it('should instrument cursor-based streaming with pg-cursor', async () => {
+                  const tracingPromise = agent.use(traces => {
+                    expect(traces[0][0]).to.have.property('name', expectedSchema.outbound.opName)
+                    expect(traces[0][0]).to.have.property('service', expectedSchema.outbound.serviceName)
+                    expect(traces[0][0]).to.have.property('resource', 'SELECT * FROM generate_series(0, 1) num')
+                    expect(traces[0][0]).to.have.property('type', 'sql')
+                    expect(traces[0][0].meta).to.have.property('span.kind', 'client')
+                    expect(traces[0][0].meta).to.have.property('db.name', 'postgres')
+                    expect(traces[0][0].meta).to.have.property('db.type', 'postgres')
+                    expect(traces[0][0].meta).to.have.property('component', 'pg')
+                    expect(traces[0][0].metrics).to.have.property('db.stream', 1)
+                    expect(traces[0][0].metrics).to.have.property('network.destination.port', 5432)
+                  })
+
+                  const cursor = client.query(new Cursor('SELECT * FROM generate_series(0, 1) num'))
+
+                  cursor.read(1, () => {
+                    cursor.close()
+                  })
+                  await tracingPromise
+                })
+              })
+
+              withVersions('pg', 'pg-query-stream', pgQueryStreamVersion => {
+                let QueryStream
+
+                beforeEach(() => {
+                  QueryStream = require(`../../../versions/pg-query-stream@${pgQueryStreamVersion}`).get()
+                })
+
+                it('should instrument stream-based queries with pg-query-stream', async () => {
+                  const agentPromise = agent.use(traces => {
+                    expect(traces[0][0]).to.have.property('name', expectedSchema.outbound.opName)
+                    expect(traces[0][0]).to.have.property('service', expectedSchema.outbound.serviceName)
+                    expect(traces[0][0]).to.have.property('resource', 'SELECT * FROM generate_series(0, 1) num')
+                    expect(traces[0][0]).to.have.property('type', 'sql')
+                    expect(traces[0][0].meta).to.have.property('span.kind', 'client')
+                    expect(traces[0][0].meta).to.have.property('db.name', 'postgres')
+                    expect(traces[0][0].meta).to.have.property('db.type', 'postgres')
+                    expect(traces[0][0].meta).to.have.property('component', 'pg')
+                    expect(traces[0][0].metrics).to.have.property('db.stream', 1)
+                    expect(traces[0][0].metrics).to.have.property('network.destination.port', 5432)
+                  })
+
+                  const query = new QueryStream('SELECT * FROM generate_series(0, 1) num', [])
+                  const stream = client.query(query)
+
+                  for await (const row of stream) {
+                    expect(row).to.have.property('num')
+                  }
+
+                  await agentPromise
+                })
+              })
+            })
+          }
         })
       })
 

packages/dd-trace/test/plugins/externals.json

@@ -431,6 +431,14 @@
     {
       "name": "express",
       "versions": [">=4"]
+    },
+    {
+      "name": "pg-query-stream",
+      "versions": [">=4"]
+    },
+    {
+      "name": "pg-cursor",
+      "versions": [">=2"]
     }
   ],
   "pino": [