fix(pg-query-stream): invoke this.callback on cursor end/error (#2810)

* fix(pg-query-stream): invoke `this.callback` on cursor end/error

Closes #2013

* fix(Client): respect `callback` argument for `Submittable` case

* Add tests

* Remove log

* Fix lint

---------

Co-authored-by: Brian Carlson <brian.m.carlson@gmail.com>

Author: aleclarson

packages/pg-query-stream/src/index.ts

@@ -13,6 +13,7 @@ class QueryStream extends Readable implements Submittable {
   cursor: any
   _result: any
 
+  callback: Function
   handleRowDescription: Function
   handleDataRow: Function
   handlePortalSuspended: Function
@@ -26,6 +27,13 @@ class QueryStream extends Readable implements Submittable {
 
     super({ objectMode: true, autoDestroy: true, highWaterMark: batchSize || highWaterMark })
     this.cursor = new Cursor(text, values, config)
+    this.cursor
+      .on('end', (result) => {
+        this.callback && this.callback(null, result)
+      })
+      .on('error', (err) => {
+        this.callback && this.callback(err)
+      })
 
     // delegate Submittable callbacks to cursor
     this.handleRowDescription = this.cursor.handleRowDescription.bind(this.cursor)

packages/pg-query-stream/test/pool.ts

@@ -0,0 +1,17 @@
+import { Pool } from 'pg'
+import QueryStream from '../src'
+
+describe('pool', function () {
+  it('works', async function () {
+    const pool = new Pool()
+    const query = new QueryStream('SELECT * FROM generate_series(0, 10) num', [])
+    const q = pool.query(query)
+    query.on('data', (row) => {
+      // just consume the whole stream
+    })
+    await q
+    query.on('end', () => {
+      pool.end()
+    })
+  })
+})

packages/pg/lib/client.js

@@ -609,8 +609,12 @@ class Client extends EventEmitter {
     } else if (typeof config.submit === 'function') {
       readTimeout = config.query_timeout || this.connectionParameters.query_timeout
       result = query = config
-      if (typeof values === 'function') {
-        query.callback = query.callback || values
+      if (!query.callback) {
+        if (typeof values === 'function') {
+          query.callback = values
+        } else if (callback) {
+          query.callback = callback
+        }
       }
     } else {
       readTimeout = config.query_timeout || this.connectionParameters.query_timeout