Improve cancel to work immediately

Depedend on the client the cancel does not work immediately.
At least the kill command works but the executing query takes some time
to receive the response from the server that the query has been
canceled.
So, once the kill command has work successfully. We cancel the promise
immediately.

src/db/clients/mysql.js

@@ -2,6 +2,8 @@ import mysql from 'mysql';
 import { identify } from 'sql-query-identifier';
 
 import createDebug from '../../debug';
+import { createCancelablePromise } from '../../utils';
+import errors from '../../errors';
 
 const debug = createDebug('db:clients:mysql');
 
@@ -178,9 +180,14 @@ export async function getTableKeys(conn, database, table) {
   }));
 }
 
+
 export function query(conn, queryText) {
   let pid = null;
   let canceling = false;
+  const cancelable = createCancelablePromise({
+    ...errors.CANCELED_BY_USER,
+    sqlectronError: 'CANCELED_BY_USER',
+  });
 
   return {
     execute() {
@@ -194,7 +201,10 @@ export function query(conn, queryText) {
         pid = dataPid[0].pid;
 
         try {
-          const data = await executeQuery(connClient, queryText);
+          const data = await Promise.race([
+            cancelable.wait(),
+            executeQuery(connClient, queryText),
+          ]);
 
           pid = null;
 
@@ -206,6 +216,8 @@ export function query(conn, queryText) {
           }
 
           throw err;
+        } finally {
+          cancelable.discard();
         }
       });
     },
@@ -220,6 +232,7 @@ export function query(conn, queryText) {
         await driverExecuteQuery(conn, {
           query: `kill ${pid};`,
         });
+        cancelable.cancel();
       } catch (err) {
         canceling = false;
         throw err;

src/db/clients/postgresql.js

@@ -2,6 +2,8 @@ import pg from 'pg';
 import { identify } from 'sql-query-identifier';
 
 import createDebug from '../../debug';
+import { createCancelablePromise } from '../../utils';
+import errors from '../../errors';
 
 const debug = createDebug('db:clients:postgresql');
 
@@ -225,6 +227,10 @@ export async function getTableKeys(conn, database, table, schema) {
 export function query(conn, queryText) {
   let pid = null;
   let canceling = false;
+  const cancelable = createCancelablePromise({
+    ...errors.CANCELED_BY_USER,
+    sqlectronError: 'CANCELED_BY_USER',
+  });
 
   return {
     execute() {
@@ -238,7 +244,10 @@ export function query(conn, queryText) {
         pid = dataPid.rows[0].pid;
 
         try {
-          const data = await executeQuery(connClient, queryText);
+          const data = await Promise.race([
+            cancelable.wait(),
+            executeQuery(connClient, queryText),
+          ]);
 
           pid = null;
 
@@ -250,6 +259,8 @@ export function query(conn, queryText) {
           }
 
           throw err;
+        } finally {
+          cancelable.discard();
         }
       });
     },
@@ -268,6 +279,8 @@ export function query(conn, queryText) {
         if (!data.rows[0].pg_cancel_backend) {
           throw new Error(`Failed canceling query with pid ${pid}.`);
         }
+
+        cancelable.cancel();
       } catch (err) {
         canceling = false;
         throw err;

src/errors.js

@@ -0,0 +1,7 @@
+export default {
+  CANCELED_BY_USER: {
+    code: 'CANCELED_BY_USER',
+    name: 'Query canceled by user',
+    message: 'Query canceled by user. The query process may still in the process list. But has already received the command to kill it successfully.',
+  },
+};

src/utils.js

@@ -71,3 +71,33 @@ export function getPort() {
     });
   });
 }
+
+export function createCancelablePromise(error, timeIdle = 100) {
+  let canceled = false;
+  let discarded = false;
+
+  const wait = (time) => new Promise((resolve) => setTimeout(resolve, time));
+
+  return {
+    async wait() {
+      while (!canceled && !discarded) {
+        await wait(timeIdle);
+      }
+
+      if (canceled) {
+        const err = new Error(error.message || 'Promise canceled.');
+
+        Object.getOwnPropertyNames(error)
+          .forEach((key) => err[key] = error[key]); // eslint-disable-line no-return-assign
+
+        throw err;
+      }
+    },
+    cancel() {
+      canceled = true;
+    },
+    discard() {
+      discarded = true;
+    },
+  };
+}