feat(vscode): gracefully handle not supported apis (#4678)

benfdking · web-flow · commit 321b118a10d2 · 2025-06-06T14:05:59.000+01:00
diff --git a/sqlmesh/lsp/custom.py b/sqlmesh/lsp/custom.py
@@ -72,3 +72,30 @@ class AllModelsForRenderResponse(PydanticModel):
     """
 
     models: t.List[ModelForRendering]
+
+
+SUPPORTED_METHODS_FEATURE = "sqlmesh/supported_methods"
+
+
+class SupportedMethodsRequest(PydanticModel):
+    """
+    Request to get all supported custom LSP methods.
+    """
+
+    pass
+
+
+class CustomMethod(PydanticModel):
+    """
+    Information about a custom LSP method.
+    """
+
+    name: str
+
+
+class SupportedMethodsResponse(PydanticModel):
+    """
+    Response containing all supported custom LSP methods.
+    """
+
+    methods: t.List[CustomMethod]
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -29,19 +29,31 @@
     ALL_MODELS_FEATURE,
     ALL_MODELS_FOR_RENDER_FEATURE,
     RENDER_MODEL_FEATURE,
+    SUPPORTED_METHODS_FEATURE,
     AllModelsRequest,
     AllModelsResponse,
     AllModelsForRenderRequest,
     AllModelsForRenderResponse,
     RenderModelRequest,
     RenderModelResponse,
+    SupportedMethodsRequest,
+    SupportedMethodsResponse,
+    CustomMethod,
 )
 from sqlmesh.lsp.hints import get_hints
 from sqlmesh.lsp.reference import get_references, get_cte_references
 from sqlmesh.lsp.uri import URI
 from web.server.api.endpoints.lineage import column_lineage, model_lineage
 from web.server.api.endpoints.models import get_models
 
+SUPPORTED_CUSTOM_METHODS = [
+    ALL_MODELS_FEATURE,
+    RENDER_MODEL_FEATURE,
+    ALL_MODELS_FOR_RENDER_FEATURE,
+    API_FEATURE,
+    SUPPORTED_METHODS_FEATURE,
+]
+
 
 class SQLMeshLanguageServer:
     def __init__(
@@ -149,6 +161,20 @@ def all_models_for_render(
                 models=self.lsp_context.list_of_models_for_rendering()
             )
 
+        @self.server.feature(SUPPORTED_METHODS_FEATURE)
+        def supported_methods(
+            ls: LanguageServer, params: SupportedMethodsRequest
+        ) -> SupportedMethodsResponse:
+            """Return all supported custom LSP methods."""
+            return SupportedMethodsResponse(
+                methods=[
+                    CustomMethod(
+                        name=name,
+                    )
+                    for name in SUPPORTED_CUSTOM_METHODS
+                ]
+            )
+
         @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/renderModel.ts b/vscode/extension/src/commands/renderModel.ts
@@ -73,7 +73,7 @@ export function renderModel(
 
       if (isErr(allModelsResult)) {
         vscode.window.showErrorMessage(
-          `Failed to get models: ${allModelsResult.error}`,
+          `Failed to get models: ${allModelsResult.error.message}`,
         )
         return
       }
@@ -115,7 +115,9 @@ export function renderModel(
     })
 
     if (isErr(result)) {
-      vscode.window.showErrorMessage(`Failed to render model: ${result.error}`)
+      vscode.window.showErrorMessage(
+        `Failed to render model: ${result.error.message}`,
+      )
       return
     }
 
diff --git a/vscode/extension/src/lsp/custom.ts b/vscode/extension/src/lsp/custom.ts
@@ -31,6 +31,7 @@ export type CustomLSPMethods =
   | AbstractAPICall
   | RenderModelMethod
   | AllModelsForRenderMethod
+  | SupportedMethodsMethod
 
 interface AllModelsRequest {
   textDocument: {
@@ -75,3 +76,20 @@ export interface ModelForRendering {
   description: string | null | undefined
   uri: string
 }
+
+export interface SupportedMethodsMethod {
+  method: 'sqlmesh/supported_methods'
+  request: SupportedMethodsRequest
+  response: SupportedMethodsResponse
+}
+
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+interface SupportedMethodsRequest {}
+
+interface SupportedMethodsResponse {
+  methods: CustomMethod[]
+}
+
+interface CustomMethod {
+  name: string
+}
diff --git a/vscode/extension/src/lsp/lsp.ts b/vscode/extension/src/lsp/lsp.ts
@@ -9,18 +9,41 @@ import { sqlmeshLspExec } from '../utilities/sqlmesh/sqlmesh'
 import { err, isErr, ok, Result } from '@bus/result'
 import { getWorkspaceFolders } from '../utilities/common/vscodeapi'
 import { traceError, traceInfo } from '../utilities/common/log'
-import { ErrorType } from '../utilities/errors'
+import {
+  ErrorType,
+  ErrorTypeGeneric,
+  ErrorTypeInvalidState,
+  ErrorTypeSQLMeshOutdated,
+} from '../utilities/errors'
 import { CustomLSPMethods } from './custom'
 
+type SupportedMethodsState =
+  | { type: 'not-fetched' }
+  | { type: 'fetched'; methods: Set<string> }
+  // TODO: This state is used when the `sqlmesh/supported_methods` endpoint is
+  //  not supported by the LSP server. This is in order to be backward compatible
+  //  with older versions of SQLMesh that do not support this endpoint. At some point
+  //  we should remove this state and always fetch the supported methods.
+  | { type: 'endpoint-not-supported' }
+
 let outputChannel: OutputChannel | undefined
 
 export class LSPClient implements Disposable {
   private client: LanguageClient | undefined
+  /**
+   * State to track whether the supported methods have been fetched. These are used to determine if a method is supported
+   * by the LSP server and return an error if not.
+   */
+  private supportedMethodsState: SupportedMethodsState = { type: 'not-fetched' }
 
   constructor() {
     this.client = undefined
   }
 
+  // TODO: This method is used to check if the LSP client has completion capability
+  //  in order to be backward compatible with older versions of SQLMesh that do not
+  //  support completion. At some point we should remove this method and always assume
+  //  that the LSP client has completion capability.
   public hasCompletionCapability(): boolean {
     if (!this.client) {
       traceError('LSP client is not initialized')
@@ -98,21 +121,123 @@ export class LSPClient implements Disposable {
     if (this.client) {
       await this.client.stop()
       this.client = undefined
+      // Reset supported methods state when the client stops
+      this.supportedMethodsState = { type: 'not-fetched' }
     }
   }
 
   public async dispose() {
     await this.stop()
   }
 
+  private async fetchSupportedMethods(): Promise<void> {
+    if (!this.client || this.supportedMethodsState.type !== 'not-fetched') {
+      return
+    }
+    try {
+      const result = await this.internal_call_custom_method(
+        'sqlmesh/supported_methods',
+        {},
+      )
+      if (isErr(result)) {
+        traceError(`Failed to fetch supported methods: ${result.error}`)
+        this.supportedMethodsState = { type: 'endpoint-not-supported' }
+        return
+      }
+      const methodNames = new Set(result.value.methods.map(m => m.name))
+      this.supportedMethodsState = { type: 'fetched', methods: methodNames }
+      traceInfo(
+        `Fetched supported methods: ${Array.from(methodNames).join(', ')}`,
+      )
+    } catch {
+      // If the supported_methods endpoint doesn't exist, mark it as not supported
+      this.supportedMethodsState = { type: 'endpoint-not-supported' }
+      traceInfo(
+        'Supported methods endpoint not available, proceeding without validation',
+      )
+    }
+  }
+
   public async call_custom_method<
+    Method extends Exclude<
+      CustomLSPMethods['method'],
+      'sqlmesh/supported_methods'
+    >,
+    Request extends Extract<CustomLSPMethods, { method: Method }>['request'],
+    Response extends Extract<CustomLSPMethods, { method: Method }>['response'],
+  >(
+    method: Method,
+    request: Request,
+  ): Promise<
+    Result<
+      Response,
+      ErrorTypeGeneric | ErrorTypeInvalidState | ErrorTypeSQLMeshOutdated
+    >
+  > {
+    if (!this.client) {
+      return err({
+        type: 'generic',
+        message: 'LSP client not ready.',
+      })
+    }
+    await this.fetchSupportedMethods()
+
+    const supportedState = this.supportedMethodsState
+    switch (supportedState.type) {
+      case 'not-fetched':
+        return err({
+          type: 'invalid_state',
+          message: 'Supported methods not fetched yet whereas they should.',
+        })
+      case 'fetched': {
+        // If we have fetched the supported methods, we can check if the method is supported
+        if (!supportedState.methods.has(method)) {
+          return err({
+            type: 'sqlmesh_outdated',
+            message: `Method '${method}' is not supported by this LSP server.`,
+          })
+        }
+        const response = await this.internal_call_custom_method(
+          method,
+          request as any,
+        )
+        if (isErr(response)) {
+          return err({
+            type: 'generic',
+            message: response.error,
+          })
+        }
+        return ok(response.value as Response)
+      }
+      case 'endpoint-not-supported': {
+        const response = await this.internal_call_custom_method(
+          method,
+          request as any,
+        )
+        if (isErr(response)) {
+          return err({
+            type: 'generic',
+            message: response.error,
+          })
+        }
+        return ok(response.value as Response)
+      }
+    }
+  }
+
+  /**
+   * Internal method to call a custom LSP method without checking if the method is supported. It is used for
+   * the class whereas the `call_custom_method` checks if the method is supported.
+   */
+  public async internal_call_custom_method<
     Method extends CustomLSPMethods['method'],
     Request extends Extract<CustomLSPMethods, { method: Method }>['request'],
     Response extends Extract<CustomLSPMethods, { method: Method }>['response'],
   >(method: Method, request: Request): Promise<Result<Response, string>> {
     if (!this.client) {
       return err('lsp client not ready')
     }
+
     try {
       const result = await this.client.sendRequest<Response>(method, request)
       return ok(result)
diff --git a/vscode/extension/src/utilities/errors.ts b/vscode/extension/src/utilities/errors.ts
@@ -7,13 +7,53 @@ import { traceInfo } from './common/log'
  * Represents different types of errors that can occur in the application.
  */
 export type ErrorType =
-  | { type: 'generic'; message: string }
+  | ErrorTypeGeneric
   | { type: 'not_signed_in' }
   | { type: 'sqlmesh_lsp_not_found' }
   // tcloud_bin_not_found is used when the tcloud executable is not found. This is likely to happen if the user
   // opens a project that has a `tcloud.yaml` file but doesn't have tcloud installed.
   | { type: 'tcloud_bin_not_found' }
   | SqlmeshLspDependenciesMissingError
+  | ErrorTypeInvalidState
+  | ErrorTypeSQLMeshOutdated
+
+/**
+ * ErrorTypeSQLMeshOutdated is used when the SQLMesh version is outdated. The
+ * message should explain the problem, but the suggestion to update SQLMesh is
+ * handled at the place where the error is shown.
+ */
+export interface ErrorTypeSQLMeshOutdated {
+  type: 'sqlmesh_outdated'
+  /**
+   * A message that describes the outdated SQLMesh version, it should not talk about
+   * updating SQLMesh. This is done at the place where the error is handled.
+   */
+  message: string
+}
+
+/**
+ * ErrorTypeInvalidState is used when the state of the application is invalid state.
+ * They should never be thrown by the application unless there is a bug in the code.
+ * The shown message should be generic and not contain any sensitive information but
+ * asks the user to report the issue to the developers.
+ */
+export interface ErrorTypeInvalidState {
+  type: 'invalid_state'
+  /**
+   * A message that describes the invalid state, it should not talk about reporting
+   * the issue to the developers. This is done at the place where the error is
+   * handled.
+   */
+  message: string
+}
+
+/**
+ * ErrorTypeGeneric is a generic error type that can be used to represent any error with a message.
+ */
+export interface ErrorTypeGeneric {
+  type: 'generic'
+  message: string
+}
 
 /**
  * SqlmeshLspDependenciesMissingError is used when the sqlmesh_lsp is found but
@@ -33,6 +73,16 @@ export async function handleError(
 ): Promise<void> {
   traceInfo('handleError', error)
   switch (error.type) {
+    case 'invalid_state':
+      await window.showErrorMessage(
+        `Invalid state: ${error.message}. Please report this issue to the developers.`,
+      )
+      return
+    case 'sqlmesh_outdated':
+      await window.showErrorMessage(
+        `SQLMesh itself is outdated. Please update SQLMesh to the latest version to use this feature. ${error.message}`,
+      )
+      return
     case 'not_signed_in':
       return handleNotSignedInError(authProvider)
     case 'sqlmesh_lsp_not_found':

Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,7 @@ export function renderModel(`
`73`	`73`
`74`	`74`	`if (isErr(allModelsResult)) {`
`75`	`75`	`vscode.window.showErrorMessage(`
`76`		- `Failed to get models: ${allModelsResult.error}`,
	`76`	+ `Failed to get models: ${allModelsResult.error.message}`,
`77`	`77`	`)`
`78`	`78`	`return`
`79`	`79`	`}`
`@@ -115,7 +115,9 @@ export function renderModel(`
`115`	`115`	`})`
`116`	`116`
`117`	`117`	`if (isErr(result)) {`
`118`		- vscode.window.showErrorMessage(`Failed to render model: ${result.error}`)
	`118`	`+ vscode.window.showErrorMessage(`
	`119`	+ `Failed to render model: ${result.error.message}`,
	`120`	`+ )`
`119`	`121`	`return`
`120`	`122`	`}`
`121`	`123`