feat(vscode): add custom api call method (#4352)

Author: benfdking

sqlmesh/lsp/api.py

@@ -0,0 +1,61 @@
+"""
+This module maps the LSP custom API calls to the SQLMesh web api.
+
+Allowing the LSP to call the web api without having to know the details of the web api
+and thus passing through the details of the web api to the LSP, so that both the LSP
+and the web api can communicate with the same process, avoiding the need to have a
+separate process for the web api.
+"""
+
+import typing as t
+from pydantic import field_validator
+from sqlmesh.utils.pydantic import PydanticModel
+from web.server.models import Model
+
+API_FEATURE = "sqlmesh/api"
+
+
+class ApiRequest(PydanticModel):
+    """
+    Request to call the SQLMesh API.
+    This is a generic request that can be used to call any API endpoint.
+    """
+
+    requestId: str
+    url: str
+    method: t.Optional[str] = "GET"
+    params: t.Optional[t.Dict[str, t.Any]] = None
+    body: t.Optional[t.Dict[str, t.Any]] = None
+
+
+class ApiResponseGetModels(PydanticModel):
+    """
+    Response from the SQLMesh API for the get_models endpoint.
+    """
+
+    data: t.List[Model]
+
+    @field_validator("data", mode="before")
+    def sanitize_datetime_fields(cls, data: t.List[Model]) -> t.List[Model]:
+        """
+        Convert datetime objects to None to avoid serialization issues.
+        """
+        if isinstance(data, list):
+            for model in data:
+                if hasattr(model, "details") and model.details:
+                    # Convert datetime fields to None to avoid serialization issues
+                    for field in ["stamp", "start", "cron_prev", "cron_next"]:
+                        if (
+                            hasattr(model.details, field)
+                            and getattr(model.details, field) is not None
+                        ):
+                            setattr(model.details, field, None)
+        return data
+
+
+class ApiResponseGetLineage(PydanticModel):
+    """
+    Response from the SQLMesh API for the get_lineage endpoint.
+    """
+
+    data: t.Dict[str, t.List[str]]

sqlmesh/lsp/main.py

@@ -11,12 +11,20 @@
 from sqlmesh._version import __version__
 from sqlmesh.core.context import Context
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
+from sqlmesh.lsp.api import (
+    API_FEATURE,
+    ApiRequest,
+    ApiResponseGetLineage,
+    ApiResponseGetModels,
+)
 from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import LSPContext, ModelTarget
 from sqlmesh.lsp.custom import ALL_MODELS_FEATURE, AllModelsRequest, AllModelsResponse
 from sqlmesh.lsp.reference import (
     get_references,
 )
+from web.server.api.endpoints.lineage import model_lineage
+from web.server.api.endpoints.models import get_models
 
 
 class SQLMeshLanguageServer:
@@ -79,6 +87,23 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
             except Exception as e:
                 return get_sql_completions(None, params.textDocument.uri)
 
+        @self.server.feature(API_FEATURE)
+        def api(
+            ls: LanguageServer, request: ApiRequest
+        ) -> t.Union[ApiResponseGetModels, ApiResponseGetLineage]:
+            ls.log_trace(f"API request: {request}")
+            if self.lsp_context is None:
+                raise RuntimeError("No context found")
+            if request.url == "/api/models":
+                response = ApiResponseGetModels(data=get_models(self.lsp_context.context))
+                return response
+            if request.url.startswith("/api/lineage"):
+                name = request.url.split("/")[-1]
+                lineage = model_lineage(name, self.lsp_context.context)
+                non_set_lineage = {k: v for k, v in lineage.items() if v is not None}
+                return ApiResponseGetLineage(data=non_set_lineage)
+            raise NotImplementedError(f"API request not implemented: {request.url}")
+
         @self.server.feature(types.TEXT_DOCUMENT_DID_OPEN)
         def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> None:
             context = self._context_get_or_load(params.text_document.uri)