Skip to content

Commit

Permalink
LSP schema completion (#159)
Browse files Browse the repository at this point in the history
Adds schema completion for the LSP.
  • Loading branch information
joacoc authored Dec 7, 2023
1 parent e1958c1 commit 97f3588
Show file tree
Hide file tree
Showing 5 changed files with 199 additions and 53 deletions.
107 changes: 74 additions & 33 deletions src/clients/lsp.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import stream from "stream";
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";

Expand Down Expand Up @@ -61,30 +62,37 @@ enum State {
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);
}
};

Expand Down Expand Up @@ -148,7 +156,7 @@ export default class LspClient {
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");
});
Expand Down Expand Up @@ -228,20 +236,21 @@ export default class LspClient {
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
Expand All @@ -257,13 +266,15 @@ export default class LspClient {
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,
}
};

Expand Down Expand Up @@ -322,10 +333,40 @@ export default class LspClient {

/**
* 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;
}
}

/**
Expand Down
97 changes: 84 additions & 13 deletions src/context/asyncContext.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
import { AdminClient, CloudClient, SqlClient } from "../clients";
import { ExtensionContext } from "vscode";
import { Context } from "./context";
import { Context, SchemaObject, SchemaObjectColumn } from "./context";

Check warning on line 3 in src/context/asyncContext.ts

View workflow job for this annotation

GitHub Actions / build (ubuntu-latest)

'SchemaObject' is defined but never used
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.
Expand Down Expand Up @@ -173,9 +174,9 @@ export default class AsyncContext extends Context {
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 {
Expand All @@ -185,32 +186,35 @@ export default class AsyncContext extends Context {
{ 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);

Check warning on line 192 in src/context/asyncContext.ts

View workflow job for this annotation

GitHub Actions / build (ubuntu-latest)

Unexpected any. Specify a different type

this.environment = {
cluster,
database,
schema,
databases,
schemas: schemas.filter(x => x.databaseId === databaseObj?.id),
schemas: schemas.filter((x: { databaseId: any; }) => x.databaseId === databaseObj?.id),

Check warning on line 199 in src/context/asyncContext.ts

View workflow job for this annotation

GitHub Actions / build (ubuntu-latest)

Unexpected any. Specify a different type
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 }] },
Expand All @@ -220,14 +224,72 @@ export default class AsyncContext extends Context {
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) => {

Check warning on line 272 in src/context/asyncContext.ts

View workflow job for this annotation

GitHub Actions / build (ubuntu-latest)

Unexpected any. Specify a different type
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.
*
Expand Down Expand Up @@ -529,4 +591,13 @@ export default class AsyncContext extends Context {
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();
}
}
40 changes: 35 additions & 5 deletions src/context/context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*/
Expand All @@ -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
Expand All @@ -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;
}
Expand Down
Loading

0 comments on commit 97f3588

Please sign in to comment.