feat(lsp): custom methods always return errors (#4708)

Author: benfdking

sqlmesh/lsp/api.py

@@ -9,13 +9,16 @@
 
 import typing as t
 from pydantic import field_validator
-from sqlmesh.utils.pydantic import PydanticModel
+from sqlmesh.lsp.custom import (
+    CustomMethodRequestBaseClass,
+    CustomMethodResponseBaseClass,
+)
 from web.server.models import LineageColumn, Model
 
 API_FEATURE = "sqlmesh/api"
 
 
-class ApiRequest(PydanticModel):
+class ApiRequest(CustomMethodRequestBaseClass):
     """
     Request to call the SQLMesh API.
     This is a generic request that can be used to call any API endpoint.
@@ -28,7 +31,11 @@ class ApiRequest(PydanticModel):
     body: t.Optional[t.Dict[str, t.Any]] = None
 
 
-class ApiResponseGetModels(PydanticModel):
+class BaseAPIResponse(CustomMethodResponseBaseClass):
+    error: t.Optional[str] = None
+
+
+class ApiResponseGetModels(BaseAPIResponse):
     """
     Response from the SQLMesh API for the get_models endpoint.
     """
@@ -53,15 +60,15 @@ def sanitize_datetime_fields(cls, data: t.List[Model]) -> t.List[Model]:
         return data
 
 
-class ApiResponseGetLineage(PydanticModel):
+class ApiResponseGetLineage(BaseAPIResponse):
     """
     Response from the SQLMesh API for the get_lineage endpoint.
     """
 
     data: t.Dict[str, t.List[str]]
 
 
-class ApiResponseGetColumnLineage(PydanticModel):
+class ApiResponseGetColumnLineage(BaseAPIResponse):
     """
     Response from the SQLMesh API for the get_column_lineage endpoint.
     """

sqlmesh/lsp/custom.py

@@ -2,18 +2,28 @@
 import typing as t
 from sqlmesh.utils.pydantic import PydanticModel
 
+
+class CustomMethodRequestBaseClass(PydanticModel):
+    pass
+
+
+class CustomMethodResponseBaseClass(PydanticModel):
+    # Prefixing, so guaranteed not to collide
+    response_error: t.Optional[str] = None
+
+
 ALL_MODELS_FEATURE = "sqlmesh/all_models"
 
 
-class AllModelsRequest(PydanticModel):
+class AllModelsRequest(CustomMethodRequestBaseClass):
     """
     Request to get all the models that are in the current project.
     """
 
     textDocument: types.TextDocumentIdentifier
 
 
-class AllModelsResponse(PydanticModel):
+class AllModelsResponse(CustomMethodResponseBaseClass):
     """
     Response to get all the models that are in the current project.
     """
@@ -26,7 +36,7 @@ class AllModelsResponse(PydanticModel):
 RENDER_MODEL_FEATURE = "sqlmesh/render_model"
 
 
-class RenderModelRequest(PydanticModel):
+class RenderModelRequest(CustomMethodRequestBaseClass):
     textDocumentUri: str
 
 
@@ -41,7 +51,7 @@ class RenderModelEntry(PydanticModel):
     rendered_query: str
 
 
-class RenderModelResponse(PydanticModel):
+class RenderModelResponse(CustomMethodResponseBaseClass):
     """
     Response to render a model.
     """
@@ -63,11 +73,11 @@ class ModelForRendering(PydanticModel):
     uri: str
 
 
-class AllModelsForRenderRequest(PydanticModel):
+class AllModelsForRenderRequest(CustomMethodRequestBaseClass):
     pass
 
 
-class AllModelsForRenderResponse(PydanticModel):
+class AllModelsForRenderResponse(CustomMethodResponseBaseClass):
     """
     Response to get all the models that are in the current project for rendering purposes.
     """
@@ -94,7 +104,7 @@ class CustomMethod(PydanticModel):
     name: str
 
 
-class SupportedMethodsResponse(PydanticModel):
+class SupportedMethodsResponse(CustomMethodResponseBaseClass):
     """
     Response containing all supported custom LSP methods.
     """
@@ -105,15 +115,15 @@ class SupportedMethodsResponse(PydanticModel):
 FORMAT_PROJECT_FEATURE = "sqlmesh/format_project"
 
 
-class FormatProjectRequest(PydanticModel):
+class FormatProjectRequest(CustomMethodRequestBaseClass):
     """
     Request to format all models in the current project.
     """
 
     pass
 
 
-class FormatProjectResponse(PydanticModel):
+class FormatProjectResponse(CustomMethodResponseBaseClass):
     """
     Response to format project request.
     """