diff --git a/sqlmesh/lsp/custom.py b/sqlmesh/lsp/custom.py index a6a2de71dc..9e0bc07cd4 100644 --- a/sqlmesh/lsp/custom.py +++ b/sqlmesh/lsp/custom.py @@ -99,3 +99,22 @@ class SupportedMethodsResponse(PydanticModel): """ methods: t.List[CustomMethod] + + +FORMAT_PROJECT_FEATURE = "sqlmesh/format_project" + + +class FormatProjectRequest(PydanticModel): + """ + Request to format all models in the current project. + """ + + pass + + +class FormatProjectResponse(PydanticModel): + """ + Response to format project request. + """ + + pass diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py index d9a34b8004..0e7d89be2a 100755 --- a/sqlmesh/lsp/main.py +++ b/sqlmesh/lsp/main.py @@ -30,6 +30,7 @@ ALL_MODELS_FOR_RENDER_FEATURE, RENDER_MODEL_FEATURE, SUPPORTED_METHODS_FEATURE, + FORMAT_PROJECT_FEATURE, AllModelsRequest, AllModelsResponse, AllModelsForRenderRequest, @@ -38,6 +39,8 @@ RenderModelResponse, SupportedMethodsRequest, SupportedMethodsResponse, + FormatProjectRequest, + FormatProjectResponse, CustomMethod, ) from sqlmesh.lsp.hints import get_hints @@ -52,6 +55,7 @@ ALL_MODELS_FOR_RENDER_FEATURE, API_FEATURE, SUPPORTED_METHODS_FEATURE, + FORMAT_PROJECT_FEATURE, ] @@ -175,6 +179,25 @@ def supported_methods( ] ) + @self.server.feature(FORMAT_PROJECT_FEATURE) + def format_project( + ls: LanguageServer, params: FormatProjectRequest + ) -> FormatProjectResponse: + """Format all models in the current project.""" + try: + if self.lsp_context is None: + current_path = Path.cwd() + self._ensure_context_in_folder(current_path) + if self.lsp_context is None: + raise RuntimeError("No context found") + + # Call the format method on the context + self.lsp_context.context.format() + return FormatProjectResponse() + except Exception as e: + ls.log_trace(f"Error formatting project: {e}") + return FormatProjectResponse() + @self.server.feature(API_FEATURE) def api(ls: LanguageServer, request: ApiRequest) -> t.Dict[str, t.Any]: ls.log_trace(f"API request: {request}") diff --git a/vscode/extension/src/commands/format.ts b/vscode/extension/src/commands/format.ts index 45ef3dae7d..5e3465921a 100644 --- a/vscode/extension/src/commands/format.ts +++ b/vscode/extension/src/commands/format.ts @@ -5,19 +5,45 @@ import * as vscode from 'vscode' import { ErrorType, handleError } from '../utilities/errors' import { AuthenticationProviderTobikoCloud } from '../auth/auth' import { execAsync } from '../utilities/exec' +import { LSPClient } from '../lsp/lsp' export const format = - (authProvider: AuthenticationProviderTobikoCloud) => + ( + authProvider: AuthenticationProviderTobikoCloud, + lsp: LSPClient | undefined, + ) => async (): Promise => { traceLog('Calling format') - const out = await internalFormat() + const out = await internalFormat(lsp) if (isErr(out)) { return handleError(authProvider, out.error, 'Project format failed') } vscode.window.showInformationMessage('Project formatted successfully') } -const internalFormat = async (): Promise> => { +const internalFormat = async ( + lsp: LSPClient | undefined, +): Promise> => { + try { + // Try LSP method first + if (lsp) { + const response = await lsp.call_custom_method( + 'sqlmesh/format_project', + {}, + ) + if (isErr(response)) { + return response + } + return ok(undefined) + } + } catch (error) { + traceLog(`LSP format failed, falling back to CLI: ${JSON.stringify(error)}`) + } + + // Fallback to CLI method if LSP is not available + // TODO This is a solution in order to be backwards compatible in the cases + // where the LSP method is not implemented yet. This should be removed at + // some point in the future. const exec = await sqlmeshExec() if (isErr(exec)) { return exec diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts index a77cfe6a9e..8749c61fb2 100644 --- a/vscode/extension/src/extension.ts +++ b/vscode/extension/src/extension.ts @@ -56,9 +56,6 @@ export async function activate(context: vscode.ExtensionContext) { context.subscriptions.push( vscode.commands.registerCommand('sqlmesh.signout', signOut(authProvider)), ) - context.subscriptions.push( - vscode.commands.registerCommand('sqlmesh.format', format(authProvider)), - ) lspClient = new LSPClient() @@ -79,6 +76,13 @@ export async function activate(context: vscode.ExtensionContext) { ), ) + context.subscriptions.push( + vscode.commands.registerCommand( + 'sqlmesh.format', + format(authProvider, lspClient), + ), + ) + // Register the webview const lineagePanel = new LineagePanel(context.extensionUri, lspClient) context.subscriptions.push( diff --git a/vscode/extension/src/lsp/custom.ts b/vscode/extension/src/lsp/custom.ts index 882209de17..be11419b79 100644 --- a/vscode/extension/src/lsp/custom.ts +++ b/vscode/extension/src/lsp/custom.ts @@ -32,6 +32,7 @@ export type CustomLSPMethods = | RenderModelMethod | AllModelsForRenderMethod | SupportedMethodsMethod + | FormatProjectMethod interface AllModelsRequest { textDocument: { @@ -93,3 +94,15 @@ interface SupportedMethodsResponse { interface CustomMethod { name: string } + +export interface FormatProjectMethod { + method: 'sqlmesh/format_project' + request: FormatProjectRequest + response: FormatProjectResponse +} + +// eslint-disable-next-line @typescript-eslint/no-empty-object-type +interface FormatProjectRequest {} + +// eslint-disable-next-line @typescript-eslint/no-empty-object-type +interface FormatProjectResponse {} diff --git a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts index 707e57f76b..95cc94d38e 100644 --- a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts +++ b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts @@ -197,7 +197,7 @@ export const ensureSqlmeshEnterpriseInstalled = async (): Promise< /** * Get the sqlmesh executable for the current workspace. * - * @returns The sqlmesh executable for the current workspace. + * @deprecated Use LSP instead of direct sqlmesh execution for any new functionality. */ export const sqlmeshExec = async (): Promise< Result diff --git a/vscode/extension/tests/format.spec.ts b/vscode/extension/tests/format.spec.ts new file mode 100644 index 0000000000..ad85340aa3 --- /dev/null +++ b/vscode/extension/tests/format.spec.ts @@ -0,0 +1,48 @@ +import { test, expect } from '@playwright/test' +import path from 'path' +import fs from 'fs-extra' +import os from 'os' +import { startVSCode, SUSHI_SOURCE_PATH } from './utils' + +test('Format project works correctly', async () => { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-')) + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + try { + const { window, close } = await startVSCode(tempDir) + + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder, excluding external_models + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the customer_revenue_lifetime model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + await window.waitForSelector('text=grain') + await window.waitForSelector('text=Loaded SQLMesh Context') + + // Render the model + await window.keyboard.press( + process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P', + ) + await window.keyboard.type('Format SQLMesh Project') + await window.keyboard.press('Enter') + + // Check that the notification appears saying 'Project formatted successfully' + await expect( + window.getByText('Project formatted successfully', { exact: true }), + ).toBeVisible() + + await close() + } finally { + await fs.remove(tempDir) + } +})