feat(vscode): move format functionality to lsp (#4681)

Author: benfdking

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

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}")

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<void> => {
     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<Result<undefined, ErrorType>> => {
+const internalFormat = async (
+  lsp: LSPClient | undefined,
+): Promise<Result<undefined, ErrorType>> => {
+  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

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(

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 {}

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<SqlmeshExecInfo, ErrorType>

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)
+  }
+})