Changing instrumentation for query object and adding tests

DataDog · Aug 2, 2024 · 8fad11b · 8fad11b
1 parent d46be5e
commit 8fad11b
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 60 deletions.
diff --git a/packages/datadog-instrumentations/src/pg.js b/packages/datadog-instrumentations/src/pg.js
@@ -72,17 +72,27 @@ function wrapQuery (query) {
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
         const reusingQuery = typeof pgQuery.submit === 'function'
-        const callback = (reusingQuery && pgQuery.callback) || arguments[arguments.length - 1]
+        const callback = arguments[arguments.length - 1]
 
         finish(error)
 
-        if (typeof callback === 'function') {
-          callback(error)
+        if (reusingQuery) {
+          if (typeof callback === 'function') {
+            pgQuery.callback = pgQuery.callback || callback
+          }
 
-          if (reusingQuery) {
-            return pgQuery
+          if (pgQuery.callback) {
+            pgQuery.callback(error)
+          } else {
+            pgQuery.emit('error', error)
           }
 
+          return pgQuery
+        }
+
+        if (typeof callback === 'function') {
+          callback(error)
+
           return
         } else {
           return Promise.reject(error)

diff --git a/packages/datadog-instrumentations/test/pg.spec.js b/packages/datadog-instrumentations/test/pg.spec.js
@@ -15,6 +15,7 @@ describe('pg instrumentation', () => {
 
     let pg
     let client
+    let Query
 
     function abortQuery ({ abortController }) {
       const error = new Error('Test')
@@ -30,75 +31,104 @@ describe('pg instrumentation', () => {
     })
 
     Object.keys(clients).forEach(implementation => {
-      beforeEach(done => {
-        pg = require(`../../../versions/pg@${version}`).get()
-        const Client = clients[implementation](pg)
-
-        client = new Client({
-          host: '127.0.0.1',
-          user: 'postgres',
-          password: 'postgres',
-          database: 'postgres',
-          application_name: 'test'
-        })
+      describe(implementation, () => {
+        beforeEach(done => {
+          pg = require(`../../../versions/pg@${version}`).get()
+          const Client = clients[implementation](pg)
+          Query = Client.Query
+          client = new Client({
+            host: '127.0.0.1',
+            user: 'postgres',
+            password: 'postgres',
+            database: 'postgres',
+            application_name: 'test'
+          })
 
-        client.connect(err => done(err))
-      })
+          client.connect(err => done(err))
+        })
 
-      // afterEach((done) => {
-      //   client.end((err) => {
-      //     done(err)
-      //   })
-      // })
-
-      describe('pg.Client', () => {
-        describe('abortController', () => {
-          afterEach(() => {
-            if (queryStartChannel.hasSubscribers) {
-              queryStartChannel.unsubscribe(abortQuery)
-            }
-          })
+        describe('pg.Client', () => {
+          describe('abortController', () => {
+            afterEach(() => {
+              if (queryStartChannel.hasSubscribers) {
+                queryStartChannel.unsubscribe(abortQuery)
+              }
+            })
 
-          describe('using callback', () => {
-            it('Should not fail if it is not aborted', (done) => {
-              client.query('SELECT 1', (err) => {
-                client.end()
-                done(err)
+            describe('using callback', () => {
+              it('Should not fail if it is not aborted', (done) => {
+                client.query('SELECT 1', (err) => {
+                  client.end()
+                  done(err)
+                })
               })
-            })
 
-            it('Should abort query', (done) => {
-              queryStartChannel.subscribe(abortQuery)
+              it('Should abort query', (done) => {
+                queryStartChannel.subscribe(abortQuery)
 
-              client.query('SELECT 1', (err) => {
-                client.end()
-                if (err && err.message === 'Test') {
-                  return done()
-                }
+                client.query('SELECT 1', (err) => {
+                  client.end()
+                  if (err && err.message === 'Test') {
+                    return done()
+                  }
 
-                done(new Error('Query was not aborted'))
+                  done(new Error('Query was not aborted'))
+                })
               })
             })
-          })
 
-          describe('using promise', () => {
-            it('Should not fail if it is not aborted', async () => {
-              await client.query('SELECT 1')
-              client.end()
-            })
+            describe('using promise', () => {
+              it('Should not fail if it is not aborted', async () => {
+                await client.query('SELECT 1')
+                client.end()
+              })
 
-            it('Should abort query', async () => {
-              queryStartChannel.subscribe(abortQuery)
+              it('Should abort query', async () => {
+                queryStartChannel.subscribe(abortQuery)
 
-              try {
-                await client.query('SELECT 1')
+                try {
+                  await client.query('SELECT 1')
 
-                throw new Error('Query was not aborted')
-              } catch (err) {
-                if (!err || err.message !== 'Test') {
-                  throw err
+                  throw new Error('Query was not aborted')
+                } catch (err) {
+                  if (!err || err.message !== 'Test') {
+                    throw err
+                  }
                 }
-              }
+              })
+            })
+
+            describe('using query object', () => {
+              describe('without callback', () => {
+                it('Should not fail if it is not aborted', (done) => {
+                  const query = new Query('SELECT 1')
+                  client.query(query)
+                  query.on('end', () => {
+                    client.end()
+                    done()
+                  })
+                })
+
+                it('Should abort query', (done) => {
+                  queryStartChannel.subscribe(abortQuery)
+                  const query = new Query('SELECT 1')
+
+                  client.query(query)
+                  query.on('error', err => {
+                    console.log('que pasa aquí?', err.message, err)
+
+                    if (err && err.message === 'Test') {
+                      done()
+                      return
+                    }
+                    done(err)
+                  })
+
+                  query.on('end', () => {
+                    done(new Error('Query was not aborted'))
+                  })
+                })
+              })
             })
           })
         })