feat(vscode): add end point for column lineage

Author: benfdking

sqlmesh/lsp/api.py

@@ -10,7 +10,7 @@
 import typing as t
 from pydantic import field_validator
 from sqlmesh.utils.pydantic import PydanticModel
-from web.server.models import Model
+from web.server.models import LineageColumn, Model
 
 API_FEATURE = "sqlmesh/api"
 
@@ -59,3 +59,11 @@ class ApiResponseGetLineage(PydanticModel):
     """
 
     data: t.Dict[str, t.List[str]]
+
+
+class ApiResponseGetColumnLineage(PydanticModel):
+    """
+    Response from the SQLMesh API for the get_column_lineage endpoint.
+    """
+
+    data: t.Dict[str, t.Dict[str, LineageColumn]]

sqlmesh/lsp/main.py

@@ -4,6 +4,7 @@
 import logging
 import typing as t
 from pathlib import Path
+import urllib.parse
 
 from lsprotocol import types
 from pygls.server import LanguageServer
@@ -14,6 +15,7 @@
 from sqlmesh.lsp.api import (
     API_FEATURE,
     ApiRequest,
+    ApiResponseGetColumnLineage,
     ApiResponseGetLineage,
     ApiResponseGetModels,
 )
@@ -24,7 +26,7 @@
     get_references,
 )
 from sqlmesh.lsp.uri import URI
-from web.server.api.endpoints.lineage import model_lineage
+from web.server.api.endpoints.lineage import column_lineage, model_lineage
 from web.server.api.endpoints.models import get_models
 
 
@@ -92,18 +94,37 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
         @self.server.feature(API_FEATURE)
         def api(
             ls: LanguageServer, request: ApiRequest
-        ) -> t.Union[ApiResponseGetModels, ApiResponseGetLineage]:
+        ) -> t.Union[ApiResponseGetModels, ApiResponseGetLineage, ApiResponseGetColumnLineage]:
             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"):
+            if (
+                request.url.startswith("/api/lineage")
+                and request.url.count("/") == 3
+                and request.method == "GET"
+            ):
                 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)
+            if (
+                request.url.startswith("/api/lineage")
+                and request.url.count("/") == 4
+                and request.method == "GET"
+            ):
+                model_name = request.url.split("/")[-2]
+                model_name = urllib.parse.unquote(model_name)
+                column = request.url.split("/")[-1]
+                column = urllib.parse.unquote(column)
+                logging.info(f"Column lineage request: {model_name} {column}")
+                column_lineage_response = column_lineage(
+                    model_name, column, False, self.lsp_context.context
+                )
+                return ApiResponseGetColumnLineage(data=column_lineage_response)
+
             raise NotImplementedError(f"API request not implemented: {request.url}")
 
         @self.server.feature(types.TEXT_DOCUMENT_DID_OPEN)

web/server/api/endpoints/lineage.py

@@ -126,7 +126,7 @@ def create_models_only_lineage_adjacency_list(
 
 
 @router.get("/{model_name:str}/{column_name:str}")
-async def column_lineage(
+def column_lineage(
     model_name: str,
     column_name: str,
     models_only: bool = False,