Add ability to cancel active query

brianc · brianc · commit 92c22c8c683d · 2020-01-28T13:55:37.000-06:00
diff --git a/packages/pg/lib/client.js b/packages/pg/lib/client.js
@@ -378,23 +378,25 @@ Client.prototype.getStartupConf = function () {
   return data
 }
 
-Client.prototype.cancel = function (client, query) {
-  if (client.activeQuery === query) {
-    var con = this.connection
+Client.prototype.cancelActiveQuery = function () {
+  var con = new Connection({
+    ssl: this.connectionParameters.ssl
+  })
 
-    if (this.host && this.host.indexOf('/') === 0) {
-      con.connect(this.host + '/.s.PGSQL.' + this.port)
-    } else {
-      con.connect(this.port, this.host)
-    }
+  if (this.host && this.host.indexOf('/') === 0) {
+    con.connect(this.host + '/.s.PGSQL.' + this.port)
+  } else {
+    con.connect(this.port, this.host)
+  }
 
+  return new Promise((resolve, reject) => {
     // once connection is established send cancel message
-    con.on('connect', function () {
-      con.cancel(client.processID, client.secretKey)
+    con.on('connect', () => {
+      con.cancel(this.processID, this.secretKey)
     })
-  } else if (client.queryQueue.indexOf(query) !== -1) {
-    client.queryQueue.splice(client.queryQueue.indexOf(query), 1)
-  }
+    con.once('end', resolve)
+    con.once('error', reject)
+  })
 }
 
 Client.prototype.setTypeParser = function (oid, format, parseFn) {
diff --git a/packages/pg/lib/native/client.js b/packages/pg/lib/native/client.js
@@ -286,6 +286,12 @@ Client.prototype.cancel = function (query) {
   }
 }
 
+Client.prototype.cancelActiveQuery = function () {
+  return new Promise((resolve, reject) => {
+    this.native.cancel((err) => err ? reject(err) : resolve())
+  })
+}
+
 Client.prototype.setTypeParser = function (oid, format, parseFn) {
   return this._types.setTypeParser(oid, format, parseFn)
 }
diff --git a/packages/pg/test/integration/client/cancel-tests.js b/packages/pg/test/integration/client/cancel-tests.js
@@ -0,0 +1,25 @@
+"use strict"
+var helper = require('./../test-helper')
+var assert = require('assert')
+
+const suite = new helper.Suite()
+
+suite.testAsync('All queries should return a result array', async () => {
+  const pool = new helper.pg.Pool()
+  const client = await pool.connect()
+  const cancledPromise = client.query('SELECT pg_sleep(100000)')
+    .then(() => {
+      throw new Error('this should not resolve because query is cancled')
+    })
+    .catch(err => {
+      return err
+    })
+  await client.cancelActiveQuery()
+  const err = await cancledPromise
+  assert(err instanceof Error)
+  // make sure client is still usable
+  const { rows } = await client.query('SELECT 1 as "foo"')
+  assert.strictEqual(rows.length, 1)
+  assert.deepStrictEqual(rows[0], { foo: 1 })
+  await client.end()
+})

Original file line number	Diff line number	Diff line change
`@@ -286,6 +286,12 @@ Client.prototype.cancel = function (query) {`
`286`	`286`	`}`
`287`	`287`	`}`
`288`	`288`
	`289`	`+Client.prototype.cancelActiveQuery = function () {`
	`290`	`+ return new Promise((resolve, reject) => {`
	`291`	`+ this.native.cancel((err) => err ? reject(err) : resolve())`
	`292`	`+ })`
	`293`	`+}`
	`294`	`+`
`289`	`295`	`Client.prototype.setTypeParser = function (oid, format, parseFn) {`
`290`	`296`	`return this._types.setTypeParser(oid, format, parseFn)`
`291`	`297`	`}`