sql: abort queries

src/clients/sql.ts

@@ -29,6 +29,36 @@
         this.handleReconnection();
     }
 
+    /**
+     * Reconnects the pool and creates a new private client.
+     *
+     * This is useful to reconnect on errors, or when wanting
+     * to abort a query.
+     */
+    async reconnect() {
+        try {
+            const client = await this.privateClient;
+            client.release();
+        } catch (err) {
+            console.error("[SqlClient]", "Error aborting private client:", err);
+        } finally {
+            try {
+                const pool = await this.pool;
+                pool.end();
+            } catch (err) {
+                console.error("[SqlClient]", "Error ending pool. It is ok it the pool connection failed:", err);
+            } finally {
+                this.pool = this.buildPool();
+                this.privateClient = this.buildPrivateClient();
+                this.handleReconnection();
+            }
+        }
+    }
+
+    /**
+     * Handles the reconnection from the pool or private client
+     * when there is a connection issue.
+     */
     private async handleReconnection() {
         let reconnecting = false;
 
@@ -38,18 +68,8 @@
             if (reconnecting === false && this.ended === false) {
                 reconnecting = true;
                 const interval = setInterval(async () => {
-                    try {
-                        const pool = await this.pool;
-                        pool.end();
-                    } catch (err) {
-                        console.error("[SqlClient]", "Error awaiting pool to end. It is ok it the pool connection failed.");
-                    } finally {
-                        this.pool = this.buildPool();
-                        this.privateClient = this.buildPrivateClient();
-                        this.handleReconnection();
-                        reconnecting = false;
-                        clearInterval(interval);
-                    }
+                    this.reconnect();
+                    clearInterval(interval);
                 }, 5000);
             }
         };
@@ -71,11 +91,17 @@
         }
     }
 
+    /**
+     * @returns a client form the pool.
+     */
     private async buildPrivateClient(): Promise<PoolClient> {
         const pool = await this.pool;
         return pool.connect();
     }
 
+    /**
+     * @returns a Postgres connection pool.
+     */
     private async buildPool(): Promise<Pool> {
         return new Pool(await this.buildPoolConfig());
     }
@@ -140,7 +166,7 @@
      * @param values
      * @returns query results
      */
     async internalQuery(statement: string, values?: Array<string | number>): Promise<QueryArrayResult<any>> {
         const pool = await this.pool;
         const results = await pool.query(statement, values);
 
@@ -161,7 +187,7 @@
      * @param values
      * @returns query results
      */
     async privateQuery(statement: string, values?: Array<any>): Promise<QueryArrayResult<any>> {
         const client = await this.privateClient;
         // Row mode is a must.
         // Otherwise when two columns have the same name, one is dropped

src/context/asyncContext.ts

@@ -1,6 +1,6 @@
 import { AdminClient, CloudClient, SqlClient } from "../clients";
 import { ExtensionContext } from "vscode";
 import { Context, SchemaObject, SchemaObjectColumn } from "./context";
 import { Errors, ExtensionError } from "../utilities/error";
 import AppPassword from "./appPassword";
 import { ActivityLogTreeProvider, AuthProvider, DatabaseTreeProvider, ResultsProvider } from "../providers";
@@ -189,14 +189,14 @@
                         { rows: schemas },
                     ] = await Promise.all(environmentPromises);
 
                     const databaseObj = databases.find((x: { name: any; }) => x.name === database);
 
                     this.environment = {
                         cluster,
                         database,
                         schema,
                         databases,
                         schemas: schemas.filter((x: { databaseId: any; }) => x.databaseId === databaseObj?.id),
                         clusters
                     };
 
@@ -269,7 +269,7 @@
         ]);
 
         const columnsMap: { [id: string] : Array<SchemaObjectColumn>; } = {};
         columnsResults.rows.forEach(({ id, name, type }: any) => {
             const columns = columnsMap[id];
             const column = { name, type };
             if (columns) {
@@ -321,6 +321,21 @@
         }
     }
 
+    /**
+     * Cancels a query by reconnecting the SQL client.
+     *
+     * This action will release the private client
+     * and create a new pool of clients.
+     *
+     */
+    async cancelQuery() {
+        try {
+            await this.clients.sql?.reconnect();
+        } catch (err) {
+            throw new ExtensionError(Errors.cancelQueryError, err);
+        }
+    }
+
     /**
      * Adds a new profile.
      *

src/providers/query.ts

@@ -33,8 +33,17 @@ export const buildRunSQLCommand = (context: AsyncContext) => {
         // the results from one query can overlap the results
         // from another. We only want to display the last results.
         const id = randomUUID();
+        let queryCanceled = false;
         try {
-            resultsProvider.setQueryId(id);
+            const cancelHandler = () => {
+                console.log("[RunSQLCommand]", "Canceling query.");
+                queryCanceled = true;
+                context.cancelQuery();
+            }
+            resultsProvider.setQuery(
+                id,
+                cancelHandler
+            );
             try {
                 const statements = await context.parseSql(query);
                 console.log("[RunSQLCommand]", "Running statements: ", statements);
@@ -78,13 +87,20 @@ export const buildRunSQLCommand = (context: AsyncContext) => {
                             sql: statement.sql
                         });
 
-                        resultsProvider.setResults(id,
-                            undefined,
-                            {
-                                message: error.toString(),
-                                position: error.position,
+                        if (!queryCanceled) {
+                            resultsProvider.setResults(id,
+                                undefined,
+                                {
+                                    message: error.toString(),
+                                    position: error.position,
+                                    query,
+                            });
+                        } else {
+                            resultsProvider.setResults(id, undefined, {
+                                message: "Query canceled.",
                                 query,
-                        });
+                            });
+                        }
 
                         // Break for-loop.
                         break;

src/providers/results.ts

@@ -21,6 +21,10 @@ export default class ResultsProvider implements vscode.WebviewViewProvider {
     // It is used to display the results and not overlap them from the results of a laggy query.
     private lastQueryId: string | undefined;
 
+    // Handles query cancelation. After running this function
+    // the context will reconnect the SQL client.
+    private cancelHandler: undefined | (() => void);
+
     // The provider can be invoked from `materialize.run`.
     // When this happens, the inner rendering script will not be ready
     // to listen changes. This variable holds the pending data to render
@@ -34,14 +38,19 @@ export default class ResultsProvider implements vscode.WebviewViewProvider {
     constructor(private readonly _extensionUri: vscode.Uri) {
         this._extensionUri = _extensionUri;
         this.isScriptReady = false;
+        this.cancelHandler = undefined;
     }
 
     /**
      * Cleans the results and sets a latest query id.
      * @param id
      */
-    public setQueryId(id: string) {
+    public setQuery(
+        id: string,
+        cancelHandler: () => void,
+    ) {
         this.lastQueryId = id;
+        this.cancelHandler = cancelHandler;
 
         if (this._view) {
             console.log("[ResultsProvider]", "New query.");
@@ -113,6 +122,12 @@ export default class ResultsProvider implements vscode.WebviewViewProvider {
                     console.error("[ResultsProvider]", error);
                     break;
                 }
+                case "cancelQuery": {
+                    if (this.cancelHandler) {
+                        this.cancelHandler();
+                    }
+                    break;
+                }
                 case "ready": {
                     console.log("[ResultsProvider]", "The script is now ready.");
                     this.isScriptReady = true;

src/providers/scripts/results.ts

@@ -78,6 +78,7 @@ import {
                         timerElement.style.paddingTop = "0.5rem";
                         document.body.appendChild(timerElement);
                     }
+
                     // Reset and show the timer content
                     timerElement.textContent = 'Time elapsed:';
                     timerElement.style.display = 'block';
@@ -131,6 +132,18 @@ import {
                             }, 500);
                         }
                     }, 103);
+
+                    const cancelQuery = () => {
+                        vscode.postMessage({ type: "cancelQuery" });
+                    };
+                    const cancelElement = document.createElement('vscode-button');
+                    cancelElement.innerHTML = "Cancel";
+                    cancelElement.style.float = "right";
+                    cancelElement.id = "cancelButton";
+                    cancelElement.setAttribute("appearance", "secondary");
+                    cancelElement.onclick = cancelQuery;
+                    container.appendChild(cancelElement);
+
                     break;
                 }
 
@@ -139,6 +152,10 @@ import {
                     if (progressRing) {
                         progressRing.style.display = "none";
                     }
+                    const cancelButton = document.getElementById("cancelButton");
+                    if (cancelButton) {
+                        cancelButton.style.display = "none";
+                    }
 
                     const { data: results } = message;
                     console.log("[Results.js]", "New message - Results: ", results);

src/utilities/error.ts

@@ -157,5 +157,9 @@ export enum Errors {
     /**
      * Raises when a fetch failes after a minute.
      */
-    fetchTimeoutError = "Failed to fetch after a minute."
+    fetchTimeoutError = "Failed to fetch after a minute.",
+    /**
+     * Raises when trying to cancel a query.
+     */
+    cancelQueryError = "Failed to cancel the query."
 }