lsp: parse before sending information

src/clients/lsp.ts

@@ -27,6 +27,21 @@ const SERVER_DECOMPRESS_PATH: string = path.join(os.tmpdir(), "mz", "bin", "mz-l
 /// The final server binary path.
 const SERVER_PATH: string = path.join(__dirname, "bin", "mz-lsp-server");
 
+
+/// Represents the structure a client uses to understand
+export interface ExecuteCommandParseStatement {
+    /// The sql content in the statement
+    sql: string,
+    /// The type of statement.
+    /// Represents the String version of [Statement].
+    kind: string,
+}
+
+/// Represents the response from the parse command.
+interface ExecuteCommandParseResponse {
+    statements: Array<ExecuteCommandParseStatement>
+}
+
 /// This class implements the Language Server Protocol (LSP) client for Materialize.
 /// The LSP is downloaded for an endpoint an it is out of the bundle. Binaries are heavy-weight
 /// and is preferable to download on the first activation.
@@ -280,4 +295,31 @@ export default class LspClient {
     stop() {
         this.client && this.client.stop();
     }
+
+    /**
+     * Sends a request to the LSP server to execute the parse command.
+     * The parse command returns the list of statements in an array,
+     * including their corresponding SQL and type (e.g., select, create_table, etc.).
+     *
+     * For more information about commands: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspace_executeCommand
+     */
+    async parseSql(sql: string): Promise<Array<ExecuteCommandParseStatement>> {
+        if (this.client) {
+            console.log("[LSP]", "Setting on request handler.");
+
+            // Setup the handler.
+            this.client.onRequest("workspace/executeCommand", (...params) => {
+                console.log("[LSP]", "Response params: ", params);
+            });
+
+            // Send request
+            const { statements } = await this.client.sendRequest("workspace/executeCommand", { command: "parse", arguments: [
+                sql
+            ]}) as ExecuteCommandParseResponse;
+
+            return statements;
+        } else {
+            throw new Error("Client is not yet available.");
+        }
+    }
 }

src/clients/sql.ts

@@ -97,6 +97,17 @@ export default class SqlClient {
         };
     }
 
+    /**
+     * Returns a client from the pool.
+     * The request must call `done()` to free the client after using it..
+     */
+    async poolClient() {
+        const pool = await this.pool;
+        const client = await pool.connect();
+
+        return client;
+    }
+
     async query(statement: string, values?: Array<any>): Promise<QueryResult<any>> {
         const pool = await this.pool;
         const results = await pool.query(statement, values);

src/context/context.ts

@@ -3,8 +3,9 @@ import { AdminClient, CloudClient, SqlClient } from "../clients";
 import { Config } from "./config";
 import { MaterializeObject, MaterializeSchemaObject } from "../providers/schema";
 import AppPassword from "./appPassword";
-import LspClient from "../clients/lsp";
+import LspClient, { ExecuteCommandParseStatement } from "../clients/lsp";
 import { Errors, ExtensionError } from "../utilities/error";
+import { PoolClient } from "pg";
 
 export enum EventType {
     newProfiles,
@@ -211,6 +212,16 @@ export class Context extends EventEmitter {
         return await client.query(text, vals);
     }
 
+    /**
+     * This method is NOT recommended to use.
+     * Make sure to understand clients from the pool lifecycle.
+     * @returns a client from the pool.
+     */
+    async poolClient(): Promise<PoolClient> {
+        const client = await this.getSqlClient();
+        return await client.poolClient();
+    }
+
     getClusters(): MaterializeObject[] | undefined {
         return this.environment?.clusters;
     }
@@ -259,6 +270,10 @@ export class Context extends EventEmitter {
         await this.loadContext();
     }
 
+    async parseSql(sql: string): Promise<Array<ExecuteCommandParseStatement>> {
+        return this.lspClient.parseSql(sql);
+    }
+
     handleErr(err: Error) {
         if (err instanceof ExtensionError) {
             this.emit("event", { type: EventType.error, message: err.message });

src/extension.ts

@@ -63,54 +63,87 @@ export function activate(vsContext: vscode.ExtensionContext) {
             const textSelected = activeEditor.document.getText(selection).trim();
             const query = textSelected ? textSelected : document.getText();
 
-            console.log("[RunSQLCommand]", "Running query: ", query);
-
             // Identify the query to not overlap results.
             // When a user press many times the run query button
             // the results from one query can overlap the results
             // from another. We only want to display the last results.
             const id = randomUUID();
-            context.emit("event", { type: EventType.newQuery, data: { id } });
 
-            // Benchmark
-            const startTime = Date.now();
             try {
-                const results = await context.query(query);
-                const endTime = Date.now();
-                const elapsedTime = endTime - startTime;
-
-                console.log("[RunSQLCommand]", "Results: ", results);
-                console.log("[RunSQLCommand]", "Emitting results.");
-
-                if (Array.isArray(results)) {
-                    context.emit("event", { type: EventType.queryResults, data: { ...results[0], elapsedTime, id } });
-                } else {
-                    context.emit("event", { type: EventType.queryResults, data: { ...results, elapsedTime, id } });
+                // Clean the results by emitting a newQuery event.
+                context.emit("event", { type: EventType.newQuery, data: { id } });
+
+                const poolClient = await context.poolClient();
+                try {
+                    const statements = await context.parseSql(query);
+
+                    console.log("[RunSQLCommand]", "Running statements: ", statements);
+
+                    const lastStatement = statements[statements.length - 1];
+                    for (const statement of statements) {
+                        console.log("[RunSQLCommand]", "Running statement: ", statement);
+
+                        // Benchmark
+                        const startTime = Date.now();
+                        try {
+                            const results = await poolClient.query(statement.sql);
+                            const endTime = Date.now();
+                            const elapsedTime = endTime - startTime;
+
+                            console.log("[RunSQLCommand]", "Results: ", results);
+                            console.log("[RunSQLCommand]", "Emitting results.");
+
+                            // Only display the results from the last statement.
+                            if (lastStatement === statement) {
+                                if (Array.isArray(results)) {
+                                    context.emit("event", { type: EventType.queryResults, data: { ...results[0], elapsedTime, id } });
+                                } else {
+                                    context.emit("event", { type: EventType.queryResults, data: { ...results, elapsedTime, id } });
+                                }
+                            }
+                            activityLogProvider.addLog({
+                                status: "success",
+                                latency: elapsedTime, // assuming elapsedTime holds the time taken for the query to execute
+                                sql: statement.sql
+                            });
+                        } catch (error: any) {
+                            console.log("[RunSQLCommand]", error.toString());
+                            console.log("[RunSQLCommand]", JSON.stringify(error));
+                            const endTime = Date.now();
+                            const elapsedTime = endTime - startTime;
+
+                            activityLogProvider.addLog({
+                                status: "failure",
+                                latency: elapsedTime, // assuming elapsedTime holds the time taken before the error was caught
+                                sql: statement.sql
+                            });
+
+                            context.emit("event", { type: EventType.queryResults, data: { id, rows: [], fields: [], error: {
+                                message: error.toString(),
+                                position: error.position,
+                                query,
+                            }, elapsedTime }});
+                            break;
+                        } finally {
+                            resultsProvider._view?.show();
+                        }
+                    }
+                } catch (err) {
+                    context.emit("event", { type: EventType.queryResults, data: { id, rows: [], fields: [], error: {
+                        message: "Syntax errors are present in your code. For more information, please check the \"Problems\" tab.",
+                        position: undefined,
+                        query,
+                    }, elapsedTime: undefined }});
+
+                    console.error("[RunSQLCommand]", "Error running statement: ", err);
+                    poolClient.release();
                 }
-                activityLogProvider.addLog({
-                    status: "success",
-                    latency: elapsedTime, // assuming elapsedTime holds the time taken for the query to execute
-                    sql: query
-                });
-            } catch (error: any) {
-                console.log("[RunSQLCommand]", error.toString());
-                console.log("[RunSQLCommand]", JSON.stringify(error));
-                const endTime = Date.now();
-                const elapsedTime = endTime - startTime;
-
-                activityLogProvider.addLog({
-                    status: "failure",
-                    latency: elapsedTime, // assuming elapsedTime holds the time taken before the error was caught
-                    sql: query
-                });
-
+            } catch (err) {
                 context.emit("event", { type: EventType.queryResults, data: { id, rows: [], fields: [], error: {
-                    message: error.toString(),
-                    position: error.position,
+                    message: "Error connecting to Materialize.",
+                    position: undefined,
                     query,
-                }, elapsedTime }});
-            } finally {
-                resultsProvider._view?.show();
+                }, elapsedTime: undefined }});
             }
         });
     });