Implement basic query cancellation for CockroachDB (#4723)

lib/dialects/cockroachdb/index.js

@@ -19,6 +19,22 @@ class Client_CockroachDB extends Client_PostgreSQL {
   _parseVersion(versionString) {
     return versionString.split(' ')[2];
   }
+
+  async cancelQuery(connectionToKill) {
+    try {
+      await this._wrappedCancelQueryCall(null, connectionToKill);
+    } catch (err) {
+      this.logger.warn(`Connection Error: ${err}`);
+      throw err;
+    }
+  }
+
+  _wrappedCancelQueryCall(emptyConnection, connectionToKill) {
+    return connectionToKill.cancel(
+      connectionToKill,
+      connectionToKill.activeQuery
+    );
+  }
 }
 
 Object.assign(Client_CockroachDB.prototype, {

lib/dialects/mysql/index.js

@@ -164,11 +164,7 @@ class Client_MySQL extends Client {
   async cancelQuery(connectionToKill) {
     const conn = await this.acquireRawConnection();
     try {
-      return await this._query(conn, {
-        sql: 'KILL QUERY ?',
-        bindings: [connectionToKill.threadId],
-        options: {},
-      });
+      return await this._wrappedCancelQueryCall(conn, connectionToKill);
     } finally {
       await this.destroyRawConnection(conn);
       if (conn.__knex__disposed) {
@@ -176,6 +172,14 @@ class Client_MySQL extends Client {
       }
     }
   }
+
+  _wrappedCancelQueryCall(conn, connectionToKill) {
+    return this._query(conn, {
+      sql: 'KILL QUERY ?',
+      bindings: [connectionToKill.threadId],
+      options: {},
+    });
+  }
 }
 
 Object.assign(Client_MySQL.prototype, {

lib/dialects/pgnative/index.js

@@ -39,6 +39,7 @@ class Client_PgNative extends Client_PG {
       return await this._wrappedCancelQueryCall(null, connectionToKill);
     } catch (err) {
       this.logger.warn(`Connection Error: ${err}`);
+      throw err;
     }
   }
 

test/integration/suite.js

@@ -32,7 +32,6 @@ module.exports = function (knex) {
     require('./query/trigger-inserts')(knex);
     require('./query/trigger-updates')(knex);
     require('./query/trigger-deletes')(knex);
-    require('./query/additional')(knex);
     require('./datatype/bigint')(knex);
     require('./datatype/decimal')(knex);
     require('./datatype/double')(knex);

test/tape/pool.js

@@ -87,7 +87,51 @@ test('#823, should not skip pool construction pool config is not defined', funct
   }
 });
 
-test('#2321 dead connections are not evicted from pool', async (t) => {
+test('#2321 dead connections are not evicted from pool (mysql)', async (t) => {
+  if (knexfile['mysql']) {
+    const knex = makeKnex(knexfile['mysql']);
+
+    t.plan(10);
+    try {
+      await Promise.all(
+        Array.from(Array(30)).map(() => {
+          // kill all connections in pool
+          return knex.raw(`KILL connection_id()`).catch(() => {
+            // just ignore errors
+          });
+        })
+      );
+
+      // TODO: It looks like this test case can trigger the race condition
+      //       outlined in this issue comment:
+      //
+      //         https://github.com/knex/knex/issues/3636#issuecomment-592005391
+      //
+      //       For now, we're working around this problem by introducing a
+      //       1 second delay.  But, we should revisit this once the connection
+      //       pool logic has been refactored.
+      await delay(1000); // wait driver to notice connection errors (2ms was enough locally)
+
+      // all connections are dead, so they should be evicted from pool and this should work
+      await Promise.all(
+        Array.from(Array(10)).map(() =>
+          knex.select(1).then(() => t.pass('Read data'))
+        )
+      );
+    } catch (e) {
+      t.fail(
+        `Should have created new connection and execute the query, got : ${e}`
+      );
+    } finally {
+      t.end();
+      await knex.destroy();
+    }
+  } else {
+    t.end();
+  }
+});
+
+test('#2321 dead connections are not evicted from pool (mysql2)', async (t) => {
   if (knexfile['mysql2']) {
     const knex = makeKnex(knexfile['mysql2']);