From 792eb6c4fd53a7bcc0ba686a02b399a4467ec665 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Thu, 7 Dec 2023 09:59:28 -0500 Subject: [PATCH 1/3] fix: reconnect on connection issues --- src/clients/sql.ts | 98 +++++++++++++++++++++++++++------------------- 1 file changed, 58 insertions(+), 40 deletions(-) diff --git a/src/clients/sql.ts b/src/clients/sql.ts index 0198fbe..a5a18e8 100644 --- a/src/clients/sql.ts +++ b/src/clients/sql.ts @@ -1,9 +1,8 @@ -import { Pool, PoolClient, PoolConfig, QueryArrayResult, QueryResult } from "pg"; +import { Pool, PoolClient, PoolConfig, QueryArrayResult } from "pg"; import AdminClient from "./admin"; import CloudClient from "./cloud"; import { Profile } from "../context/config"; import AsyncContext from "../context/asyncContext"; -import { Errors, ExtensionError } from "../utilities/error"; export default class SqlClient { private pool: Promise; @@ -12,6 +11,7 @@ export default class SqlClient { private cloudClient: CloudClient; private context: AsyncContext; private profile: Profile; + private ended: boolean; constructor( adminClient: AdminClient, @@ -23,45 +23,61 @@ export default class SqlClient { this.cloudClient = cloudClient; this.profile = profile; this.context = context; + this.ended = false; + this.pool = this.buildPool(); + this.privateClient = this.buildPrivateClient(); + this.handleReconnection(); + } - this.pool = new Promise((res, rej) => { - const asyncOp = async () => { - try { - console.log("[SqlClient]", "Building config."); - const config = await this.buildPoolConfig(); - const pool = new Pool(config); - console.log("[SqlClient]", "Connecting pool."); - - pool.connect().then(() => { - console.log("[SqlClient]", "Pool successfully connected."); - res(pool); - }).catch((err) => { - console.error(err); - rej(new ExtensionError(Errors.poolConnectionFailure, err)); - }); - } catch (err) { - console.error("[SqlClient]", "Error creating pool: ", err); - rej(new ExtensionError(Errors.poolCreationFailure, err)); - } - }; - - asyncOp(); - }); + private async handleReconnection() { + let reconnecting = false; + + const reconnect = (err: Error) => { + console.error("[SqlClient]", "Unexpected error: ", err); + console.log("[SqlClient]", "Reconnecting."); + 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); + } + }, 5000); + } + }; - this.privateClient = new Promise((res, rej) => { - const asyncOp = async () => { - try { - const pool = await this.pool; - const client = await pool.connect(); - res(client); - } catch (err) { - console.error("[SqlClient]", "Error awaiting the pool: ", err); - rej(err); - } - }; - - asyncOp(); - }); + try { + const pool = await this.pool; + pool.on("error", reconnect); + + try { + const client = await this.privateClient; + client.on("error", reconnect); + } catch (err) { + reconnect(err as Error); + console.error("[SqlClient]", "Unexpected error on client: ", err); + } + } catch (err) { + reconnect(err as Error); + console.error("[SqlClient]", "Unexpected error on pool: ", err); + } + } + + private async buildPrivateClient(): Promise { + const pool = await this.pool; + return pool.connect(); + } + + private async buildPool(): Promise { + return new Pool(await this.buildPoolConfig()); } async connectErr() { @@ -124,7 +140,7 @@ export default class SqlClient { * @param values * @returns query results */ - async internalQuery(statement: string, values?: Array): Promise> { + async internalQuery(statement: string, values?: Array): Promise> { const pool = await this.pool; const results = await pool.query(statement, values); @@ -163,6 +179,8 @@ export default class SqlClient { * Shut down cleanly the pool. */ async end() { + this.ended = true; + try { const pool = await this.pool; await pool.end(); From 43fa8cf7c54f439eef297b2f96e051c79b9609d5 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Thu, 7 Dec 2023 10:24:55 -0500 Subject: [PATCH 2/3] sql: make reconnection pub --- src/clients/sql.ts | 50 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 12 deletions(-) diff --git a/src/clients/sql.ts b/src/clients/sql.ts index a5a18e8..97db5ae 100644 --- a/src/clients/sql.ts +++ b/src/clients/sql.ts @@ -29,6 +29,36 @@ export default class SqlClient { 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 @@ export default class SqlClient { 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 @@ export default class SqlClient { } } + /** + * @returns a client form the pool. + */ private async buildPrivateClient(): Promise { const pool = await this.pool; return pool.connect(); } + /** + * @returns a Postgres connection pool. + */ private async buildPool(): Promise { return new Pool(await this.buildPoolConfig()); } From cb76caefed8e77bc03d10fa954ba1406a4e46ed8 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 13 Dec 2023 12:37:43 +0100 Subject: [PATCH 3/3] sql: abort queries --- src/context/asyncContext.ts | 15 +++++++++++++++ src/providers/query.ts | 30 +++++++++++++++++++++++------- src/providers/results.ts | 17 ++++++++++++++++- src/providers/scripts/results.ts | 17 +++++++++++++++++ src/utilities/error.ts | 6 +++++- 5 files changed, 76 insertions(+), 9 deletions(-) diff --git a/src/context/asyncContext.ts b/src/context/asyncContext.ts index 05e70df..2b0cbe0 100644 --- a/src/context/asyncContext.ts +++ b/src/context/asyncContext.ts @@ -321,6 +321,21 @@ export default class AsyncContext extends Context { } } + /** + * 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. * diff --git a/src/providers/query.ts b/src/providers/query.ts index bf42a76..d1bc88b 100644 --- a/src/providers/query.ts +++ b/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; diff --git a/src/providers/results.ts b/src/providers/results.ts index 88187c7..cb12ad8 100644 --- a/src/providers/results.ts +++ b/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; diff --git a/src/providers/scripts/results.ts b/src/providers/scripts/results.ts index 3d011f8..5b1ffed 100644 --- a/src/providers/scripts/results.ts +++ b/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); diff --git a/src/utilities/error.ts b/src/utilities/error.ts index d875513..9a3a58d 100644 --- a/src/utilities/error.ts +++ b/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." }