LSP schema completion

src/clients/lsp.ts

@@ -14,6 +14,7 @@
 import os from "os";
 import { SemVer } from "semver";
 import { Errors, ExtensionError } from "../utilities/error";
+import { ExplorerSchema } from "../context/context";
 import * as Sentry from "@sentry/node";
 import { fetchWithRetry } from "../utilities/utils";
 
@@ -61,30 +62,37 @@
 export default class LspClient {
     private isReady: Promise<boolean>;
     private client: LanguageClient | undefined;
+    private schema?: ExplorerSchema;
 
-    constructor() {
+    constructor(schema?: ExplorerSchema) {
+        this.schema = schema;
         this.isReady = new Promise((res, rej) => {
             const asyncOp = async () => {
-                if (this.isValidOs()) {
-                    if (!this.isInstalled()) {
-                        try {
-                            await this.installAndStartLspServer();
+                try {
+                    if (this.isValidOs()) {
+                        if (!this.isInstalled()) {
+                            try {
+                                await this.installAndStartLspServer();
+                                res(true);
+                            } catch (err) {
+                                Sentry.captureException(err);
+                                rej(err);
+                            }
+                        } else {
+                            console.log("[LSP]", "The server already exists.");
+                            await this.startClient();
+                            this.serverUpgradeIfAvailable();
                             res(true);
-                        } catch (err) {
-                            Sentry.captureException(err);
-                            rej(err);
                         }
                     } else {
-                        console.log("[LSP]", "The server already exists.");
-                        await this.startClient();
-                        this.serverUpgradeIfAvailable();
-                        res(true);
+                        console.error("[LSP]", "Invalid operating system.");
+                        rej(new ExtensionError(Errors.invalidOS, "Invalid operating system."));
+                        Sentry.captureException(new ExtensionError(Errors.invalidOS, "Invalid operating system."));
+                        return;
                     }
-                } else {
-                    console.error("[LSP]", "Invalid operating system.");
-                    rej(new ExtensionError(Errors.invalidOS, "Invalid operating system."));
-                    Sentry.captureException(new ExtensionError(Errors.invalidOS, "Invalid operating system."));
-                    return;
+                } catch (err) {
+                    rej(new ExtensionError(Errors.lspOnReadyFailure,err));
+                    Sentry.captureException(err);
                 }
             };
 
@@ -148,7 +156,7 @@
                     console.log("[LSP]", "Server installed.");
                     res("");
                 })
-                .on('error', (error: any) => {
+                .on('error', (error: Error) => {
                     console.error("[LSP]", "Error during decompression:", error);
                     rej("Error during compression");
                 });
@@ -228,20 +236,21 @@
         vscode.workspace.onDidChangeConfiguration(e => {
             if (e.affectsConfiguration('materialize.formattingWidth')) {
                 console.log("[LSP]", "Formatting width has changed.");
-
-                // Restart client.
-                if (this.client) {
-                    this.client.onReady().then(() => {
-                        this.stop();
-                        this.startClient();
-                    }).catch(() => {
-                        console.error("[LSP]", "Error restarting client.");
-                    });
-                }
+                this.updateServerOptions();
             }
         });
     }
 
+    /**
+     * Updates the schema and restarts the client.
+     * @param schema
+     */
+    async updateSchema(schema: ExplorerSchema) {
+        console.log("[LSP]", "Updating schema.");
+        this.schema = schema;
+        await this.updateServerOptions();
+    }
+
     /**
      * Starts the LSP Client and checks for upgrades.
      * @param serverPath
@@ -257,13 +266,15 @@
             run,
             debug: run,
         };
-        const configuration = vscode.workspace.getConfiguration('materialize');
-        const formattingWidth = configuration.get('formattingWidth');
+        const formattingWidth = this.getFormattingWidth();
         console.log("[LSP]", "Formatting width: ", formattingWidth);
+        console.log("[LSP]", "Schema: ", this.schema);
+
         const clientOptions: LanguageClientOptions = {
             documentSelector: [{ scheme: "file", language: "mzsql"}],
             initializationOptions: {
                 formattingWidth,
+                schema: this.schema,
             }
         };
 
@@ -322,10 +333,40 @@
 
     /**
      * Stops the LSP server client.
-     * This is useful before installing an upgrade.
+     *
+     * Note: This action should be executed only when deactivating the extension.
+     */
+    async stop() {
+        if (this.client) {
+            await this.client.onReady();
+            await this.client.stop();
+        }
+    }
+
+    /**
+     * @returns the current workspace formatting width.
      */
-    stop() {
-        this.client && this.client.stop();
+    private getFormattingWidth() {
+        const configuration = vscode.workspace.getConfiguration('materialize');
+        const formattingWidth = configuration.get('formattingWidth');
+
+        return formattingWidth;
+    }
+
+    /**
+     * Updates the LSP options.
+     */
+    private async updateServerOptions() {
+        // Send request
+        if (this.client) {
+            console.log("[LSP]", "Updating the server configuration.");
+            await this.client.sendRequest("workspace/executeCommand", {
+                command: "optionsUpdate",
+                arguments: [{
+                    formattingWidth: this.getFormattingWidth(),
+                    schema: this.schema,
+            }]}) as ExecuteCommandParseResponse;
+        }
     }
 
     /**
@@ -373,7 +414,7 @@
             }
         } catch (err) {
             const parsingError = "Error parsing the statements.";
             console.error("[LSP]", (err && (err as any).message));
             if (parsingError === (err instanceof Error && err.message)) {
                 throw new ExtensionError(Errors.parsingFailure, "Syntax errors present in your query.");
             } else {

src/context/asyncContext.ts

@@ -1,12 +1,13 @@
 import { AdminClient, CloudClient, SqlClient } from "../clients";
 import { ExtensionContext } from "vscode";
-import { Context } from "./context";
+import { Context, SchemaObject, SchemaObjectColumn } from "./context";
 import { Errors, ExtensionError } from "../utilities/error";
 import AppPassword from "./appPassword";
 import { ActivityLogTreeProvider, AuthProvider, DatabaseTreeProvider, ResultsProvider } from "../providers";
 import * as vscode from 'vscode';
 import { QueryArrayResult, QueryResult } from "pg";
 import { ExecuteCommandParseStatement } from "../clients/lsp";
+import { MaterializeObject } from "../providers/schema";
 
 /**
  * Represents the different providers available in the extension.
@@ -173,9 +174,9 @@
                     this.internalQuery("SHOW CLUSTER;"),
                     this.internalQuery("SHOW DATABASE;"),
                     this.internalQuery("SHOW SCHEMA;"),
-                    this.internalQuery(`SELECT id, name, owner_id as "ownerId" FROM mz_clusters;`),
-                    this.internalQuery(`SELECT id, name, owner_id as "ownerId" FROM mz_databases;`),
-                    this.internalQuery(`SELECT id, name, database_id as "databaseId", owner_id as "ownerId" FROM mz_schemas`),
+                    this.internalQuery(`SELECT id, name FROM mz_clusters;`),
+                    this.internalQuery(`SELECT id, name FROM mz_databases;`),
+                    this.internalQuery(`SELECT id, name, database_id as "databaseId" FROM mz_schemas`),
                 ];
 
                 try {
@@ -185,32 +186,35 @@
                         { rows: [{ schema }] },
                         { rows: clusters },
                         { rows: databases },
-                        { rows: schemas }
+                        { rows: schemas },
                     ] = await Promise.all(environmentPromises);
 
-                    const databaseObj = databases.find(x => x.name === database);
+                    const databaseObj = databases.find((x: { name: any; }) => x.name === database);
 
                     this.environment = {
                         cluster,
                         database,
                         schema,
                         databases,
-                        schemas: schemas.filter(x => x.databaseId === databaseObj?.id),
+                        schemas: schemas.filter((x: { databaseId: any; }) => x.databaseId === databaseObj?.id),
                         clusters
                     };
 
+                    const schemaObj = schemas.find((x: { name: string, databaseId: string, }) => x.name === schema && x.databaseId === databaseObj?.id);
+                    console.log("[AsyncContext]", schemaObj, schemas);
+                    if (schemaObj) {
+                        this.explorerSchema = await this.getExplorerSchema(database, schemaObj);
+                    }
                     console.log("[AsyncContext]", "Environment:", this.environment);
                 } catch (err) {
-                    console.error("[AsyncContext]", "Error querying evnrionment information.");
+                    console.error("[AsyncContext]", "Error querying environment information: ", err);
                     throw err;
                 }
-            }
-
-            if (reloadSchema && this.environment) {
+            } else if (reloadSchema && this.environment) {
                 console.log("[AsyncContext]", "Reloading schema.");
                 const schemaPromises = [
                     this.internalQuery("SHOW SCHEMA;"),
-                    this.internalQuery(`SELECT id, name, database_id as "databaseId", owner_id as "ownerId" FROM mz_schemas`)
+                    this.internalQuery(`SELECT id, name, database_id as "databaseId" FROM mz_schemas;`)
                 ];
                 const [
                     { rows: [{ schema }] },
@@ -220,14 +224,72 @@
                 const { databases, database } = this.environment;
                 const databaseObj = databases.find(x => x.name === database);
                 this.environment.schema = schema;
-                this.environment.schemas = schemas.filter(x => x.databaseId === databaseObj?.id);
+                this.environment.schemas = schemas.filter((x: { databaseId: string | undefined; }) => x.databaseId === databaseObj?.id);
+
+                const schemaObj = schemas.find((x: { name: string, databaseId: string, }) => x.name === schema && x.databaseId === databaseObj?.id);
+                if (schemaObj) {
+                    this.explorerSchema = await this.getExplorerSchema(database, schema);
+                }
             }
+
+            if (this.explorerSchema) {
+                console.log("[AsyncContext]", "Update schema.");
+                try {
+                    this.clients.lsp.updateSchema(this.explorerSchema);
+                } catch (err) {
+                    console.error("[AsyncContext]", "Error updating LSP schema:", err);
+                }
+            }
+
             console.log("[AsyncContext]", "Environment loaded.");
             this.loaded = true;
             return true;
         }
     }
 
+    private async getExplorerSchema(
+        database: string,
+        { name: schema, id: schemaId }: MaterializeObject,
+    ) {
+        // Not super efficient.
+        // TODO: Replace query that appears down.
+        const [columnsResults, objects] = await Promise.all([
+            this.internalQuery(`
+                SELECT * FROM mz_columns;
+            `, []),
+            this.internalQuery(`
+                SELECT id, name, 'source' AS type FROM mz_sources WHERE schema_id = $1
+                UNION ALL SELECT id, name, 'sink' AS type FROM mz_sinks WHERE schema_id = $1
+                UNION ALL SELECT id, name, 'view' AS type FROM mz_views WHERE schema_id = $1
+                UNION ALL
+                    SELECT id, name, 'materializedView' AS type FROM mz_materialized_views WHERE schema_id = $1
+                UNION ALL SELECT id, name, 'table' AS type FROM mz_tables WHERE schema_id = $1
+                ORDER BY name;
+            `, [schemaId]),
+        ]);
+
+        const columnsMap: { [id: string] : Array<SchemaObjectColumn>; } = {};
+        columnsResults.rows.forEach(({ id, name, type }: any) => {
+            const columns = columnsMap[id];
+            const column = { name, type };
+            if (columns) {
+                columns.push(column);
+            } else {
+                columnsMap[id] = [column];
+            }
+        });
+
+        return {
+            database,
+            schema,
+            objects: objects.rows.filter(x => columnsMap[x.id]).map((x: any) => ({
+                name: x.name,
+                type: x.type,
+                columns: columnsMap[x.id]
+            }))
+        };
+    }
+
     /**
      * Returns or create the SQL client once is ready.
      *
@@ -529,4 +591,13 @@
     getProviders(): Providers {
         return this.providers;
     }
+
+    /**
+     * Stops the LSP client.
+     *
+     * Note: This action should be executed only when deactivating the extension.
+     */
+    async stop() {
+        await this.clients.lsp.stop();
+    }
 }

src/context/context.ts

@@ -16,6 +16,34 @@ export interface Environment {
     cluster: string;
 }
 
+/**
+ * Represents a column from an object schema.
+ */
+export interface SchemaObjectColumn {
+    name: string;
+    type: string;
+}
+
+/**
+ * Reperesents an object from a schema.
+ *
+ * E.g. Materialized view, index.
+ */
+export interface SchemaObject {
+    type: string,
+    name: string,
+    columns: Array<SchemaObjectColumn>
+}
+
+/**
+ * Current explorer schema information
+ */
+export interface ExplorerSchema {
+    schema: string,
+    database: string,
+    objects: Array<SchemaObject>
+}
+
 /**
  * Represents the different clients available in the extension.
  */
@@ -27,7 +55,10 @@ interface Clients {
 }
 
 export class Context {
+    // Configuration file
     protected config: Config;
+
+    // Has environment finished loading?
     protected loaded: boolean;
 
     // Visual Studio Code Context
@@ -39,19 +70,18 @@ export class Context {
     // User environment
     protected environment?: Environment;
 
+    // Current exploring schema
+    protected explorerSchema?: ExplorerSchema;
+
     constructor(vsContext: vscode.ExtensionContext) {
         this.vsContext = vsContext;
         this.config = new Config();
         this.loaded = false;
         this.clients = {
-            lsp: new LspClient()
+            lsp: new LspClient(this.explorerSchema)
         };
     }
 
-    stop() {
-        this.clients.lsp.stop();
-    }
-
     isLoading(): boolean {
         return !this.loaded;
     }