diff --git a/index.js b/index.js index a3d67a6..e4c21e6 100644 --- a/index.js +++ b/index.js @@ -3,7 +3,7 @@ const fp = require('fastify-plugin') var pg = require('pg') -function transactionUtil (pool, query, values, cb) { +function transactionUtil (pool, fn, cb) { pool.connect((err, client, done) => { if (err) return cb(err) @@ -18,9 +18,8 @@ function transactionUtil (pool, query, values, cb) { client.query('BEGIN', (err) => { if (shouldAbort(err)) return cb(err) - client.query(query, values, (err, res) => { - if (shouldAbort(err)) return cb(err) + fn(client).then(res => { client.query('COMMIT', (err) => { done() if (err) { @@ -28,22 +27,24 @@ function transactionUtil (pool, query, values, cb) { } return cb(null, res) }) + }).catch(err => { + if (shouldAbort(err)) return cb(err) }) }) }) } -function transact (query, values, cb) { +function transact (fn, cb) { if (!cb) { return new Promise((resolve, reject) => { - transactionUtil(this, query, values, function (err, res) { + transactionUtil(this, fn, function (err, res) { if (err) { return reject(err) } return resolve(res) }) }) } - return transactionUtil(this, query, values, cb) + return transactionUtil(this, fn, cb) } function fastifyPostgres (fastify, options, next) { diff --git a/test.js b/test.js index 44f306c..0f4bc3d 100644 --- a/test.js +++ b/test.js @@ -268,9 +268,9 @@ test('fastify.pg.test use transact util with promise', t => { t.error(err) fastify.pg.test .query('CREATE TABLE users(id serial PRIMARY KEY, username VARCHAR (50) NOT NULL)') - .then(result => { + .then(() => { fastify.pg.test - .transact('INSERT INTO users(username) VALUES($1) RETURNING id', ['brianc']) + .transact(client => { return client.query('INSERT INTO users(username) VALUES($1) RETURNING id', ['brianc']) }) .then(result => { t.ok(result.rows[0].id === 1) fastify.pg.test @@ -308,9 +308,9 @@ test('fastify.pg.test use transact util with callback', t => { t.error(err) fastify.pg.test .query('CREATE TABLE users2(id serial PRIMARY KEY, username VARCHAR (50) NOT NULL)') - .then(function () { + .then(() => { fastify.pg.test - .transact('INSERT INTO users2(username) VALUES($1) RETURNING id', ['brianc'], function (err, res) { + .transact(client => { return client.query('INSERT INTO users2(username) VALUES($1) RETURNING id', ['brianc']) }, function (err, res) { if (err) { t.fail(err) fastify.close()