feat(vscode): ability to render model

[ci skip]

Author: benfdking

sqlmesh/lsp/context.py

@@ -3,6 +3,10 @@
 from sqlmesh.core.context import Context
 import typing as t
 
+from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.lsp.custom import RenderModelEntry
+from sqlmesh.lsp.uri import URI
+
 
 @dataclass
 class ModelTarget:
@@ -49,3 +53,18 @@ def __init__(self, context: Context) -> None:
             **model_map,
             **audit_map,
         }
+
+
+def render_model(context: LSPContext, uri: URI) -> t.Iterator[RenderModelEntry]:
+    model_target = context.map[uri.to_path()]
+    if isinstance(model_target, ModelTarget):
+        for name in model_target.names:
+            model = context.context.get_model(name)
+            if isinstance(model, SqlModel):
+                rendered_query = context.context.render(model).sql(dialect=model.dialect)
+                yield RenderModelEntry(
+                    name=model.name,
+                    fqn=model.fqn,
+                    description=model.description,
+                    rendered_query=rendered_query,
+                )

sqlmesh/lsp/custom.py

@@ -20,3 +20,29 @@ class AllModelsResponse(PydanticModel):
 
     models: t.List[str]
     keywords: t.List[str]
+
+
+RENDER_MODEL_FEATURE = "sqlmesh/render_model"
+
+
+class RenderModelRequest(PydanticModel):
+    textDocumentUri: str
+
+
+class RenderModelEntry(PydanticModel):
+    """
+    An entry in the rendered model.
+    """
+
+    name: str
+    fqn: str
+    description: str
+    rendered_query: str
+
+
+class RenderModelResponse(PydanticModel):
+    """
+    Response to render a model.
+    """
+
+    models: t.List[RenderModelEntry]

sqlmesh/lsp/main.py

@@ -20,8 +20,15 @@
     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.context import LSPContext, ModelTarget, render_model as render_model_context
+from sqlmesh.lsp.custom import (
+    ALL_MODELS_FEATURE,
+    RENDER_MODEL_FEATURE,
+    AllModelsRequest,
+    AllModelsResponse,
+    RenderModelRequest,
+    RenderModelResponse,
+)
 from sqlmesh.lsp.reference import (
     get_references,
 )
@@ -88,6 +95,12 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
             except Exception as e:
                 return get_sql_completions(None, uri)
 
+        @self.server.feature(RENDER_MODEL_FEATURE)
+        def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:
+            uri = URI(params.textDocumentUri)
+            context = self._context_get_or_load(uri)
+            return RenderModelResponse(models=list(render_model_context(context, uri)))
+
         @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/openTestQuery.ts

@@ -4,19 +4,19 @@ export function openTestQuery() {
   return async () => {
     // Get the current active editor
     const activeEditor = vscode.window.activeTextEditor
-    
+
     // Create a new untitled document with SQL language
     const document = await vscode.workspace.openTextDocument({
       language: 'sql',
-      content: 'SELECT * FROM test'
+      content: 'SELECT * FROM test',
     })
 
     // Determine the view column for side-by-side display
     let viewColumn: vscode.ViewColumn
     if (activeEditor) {
       // Open beside the current editor
-      viewColumn = activeEditor.viewColumn 
-        ? activeEditor.viewColumn + 1 
+      viewColumn = activeEditor.viewColumn
+        ? activeEditor.viewColumn + 1
         : vscode.ViewColumn.Two
     } else {
       // If no active editor, open in column two
@@ -27,7 +27,7 @@ export function openTestQuery() {
     await vscode.window.showTextDocument(document, {
       viewColumn: viewColumn,
       preview: true,
-      preserveFocus: false
+      preserveFocus: false,
     })
   }
-}
+}

vscode/extension/src/lsp/custom.ts

@@ -4,8 +4,32 @@ export interface AllModelsMethod {
   response: AllModelsResponse
 }
 
+export interface RenderModelMethod {
+  method: 'sqlmesh/render_model'
+  request: RenderModelRequest
+  response: RenderModelResponse
+}
+
+interface RenderModelRequest {
+  textDocumentUri: string
+}
+
+interface RenderModelResponse {
+  models: RenderModelEntry[]
+}
+
+interface RenderModelEntry {
+  name: string
+  fqn: string
+  description: string
+  rendered_query: string
+}
+
 // @eslint-disable-next-line  @typescript-eslint/consistent-type-definition
-export type CustomLSPMethods = AllModelsMethod | AbstractAPICall
+export type CustomLSPMethods =
+  | AllModelsMethod
+  | AbstractAPICall
+  | RenderModelMethod
 
 interface AllModelsRequest {
   textDocument: {