diff --git a/packages/pg-cursor/index.js b/packages/pg-cursor/index.js index 8e8552be8..b77fd5977 100644 --- a/packages/pg-cursor/index.js +++ b/packages/pg-cursor/index.js @@ -17,6 +17,7 @@ class Cursor extends EventEmitter { this._queue = [] this.state = 'initialized' this._result = new Result(this._conf.rowMode, this._conf.types) + this._Promise = this._conf.Promise || global.Promise this._cb = null this._rows = null this._portal = null @@ -198,6 +199,14 @@ class Cursor extends EventEmitter { } close(cb) { + let promise + + if (!cb) { + promise = new this._Promise((resolve, reject) => { + cb = (err) => (err ? reject(err) : resolve()) + }) + } + if (!this.connection || this.state === 'done') { if (cb) { return setImmediate(cb) @@ -213,23 +222,34 @@ class Cursor extends EventEmitter { cb() }) } + + // Return the promise (or undefined) + return promise } read(rows, cb) { - if (this.state === 'idle' || this.state === 'submitted') { - return this._getRows(rows, cb) - } - if (this.state === 'busy' || this.state === 'initialized') { - return this._queue.push([rows, cb]) - } - if (this.state === 'error') { - return setImmediate(() => cb(this._error)) + let promise + + if (!cb) { + promise = new this._Promise((resolve, reject) => { + cb = (err, rows) => (err ? reject(err) : resolve(rows)) + }) } - if (this.state === 'done') { - return setImmediate(() => cb(null, [])) + + if (this.state === 'idle' || this.state === 'submitted') { + this._getRows(rows, cb) + } else if (this.state === 'busy' || this.state === 'initialized') { + this._queue.push([rows, cb]) + } else if (this.state === 'error') { + setImmediate(() => cb(this._error)) + } else if (this.state === 'done') { + setImmediate(() => cb(null, [])) } else { throw new Error('Unknown state: ' + this.state) } + + // Return the promise (or undefined) + return promise } } diff --git a/packages/pg-cursor/test/promises.js b/packages/pg-cursor/test/promises.js new file mode 100644 index 000000000..7b36dab8f --- /dev/null +++ b/packages/pg-cursor/test/promises.js @@ -0,0 +1,51 @@ +const assert = require('assert') +const Cursor = require('../') +const pg = require('pg') + +const text = 'SELECT generate_series as num FROM generate_series(0, 5)' + +describe('cursor using promises', function () { + beforeEach(function (done) { + const client = (this.client = new pg.Client()) + client.connect(done) + + this.pgCursor = function (text, values) { + return client.query(new Cursor(text, values || [])) + } + }) + + afterEach(function () { + this.client.end() + }) + + it('resolve with result', async function () { + const cursor = this.pgCursor(text) + const res = await cursor.read(6) + assert.strictEqual(res.length, 6) + }) + + it('reject with error', function (done) { + const cursor = this.pgCursor('select asdfasdf') + cursor.read(1).error((err) => { + assert(err) + done() + }) + }) + + it('read multiple times', async function () { + const cursor = this.pgCursor(text) + let res + + res = await cursor.read(2) + assert.strictEqual(res.length, 2) + + res = await cursor.read(3) + assert.strictEqual(res.length, 3) + + res = await cursor.read(1) + assert.strictEqual(res.length, 1) + + res = await cursor.read(1) + assert.strictEqual(res.length, 0) + }) +})