From 8fad11b23830c1ab435d41bb3ee990042a5c0d85 Mon Sep 17 00:00:00 2001 From: Ugaitz Urien Date: Fri, 2 Aug 2024 11:05:34 +0200 Subject: [PATCH] Changing instrumentation for query object and adding tests --- packages/datadog-instrumentations/src/pg.js | 20 ++- .../datadog-instrumentations/test/pg.spec.js | 140 +++++++++++------- 2 files changed, 100 insertions(+), 60 deletions(-) diff --git a/packages/datadog-instrumentations/src/pg.js b/packages/datadog-instrumentations/src/pg.js index 5358c6e0c5b..4f4376138c4 100644 --- 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 index a54439f0e75..e72502281f4 100644 --- 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')) + }) + }) + }) }) }) })