Add support for using promises in Cursor methods (#2554)

* Add similar promise variables to read() and close() as seen in query()

* Add testing for promise specific usage

* Simplify tests as no real callbacks are involved

Removes usage of `done()` since we can end the test when we exit the function

Co-Authored-By: Charmander <[email protected]>

* Switch to let over var

Co-authored-by: Charmander <[email protected]>

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
   }
 }
 

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)
+  })
+})