feat(lsp): include model completion metadata

Author: benfdking

sqlmesh/lsp/completions.py

@@ -1,10 +1,13 @@
 from functools import lru_cache
 from sqlglot import Dialect, Tokenizer
-from sqlmesh.lsp.custom import AllModelsResponse, MacroCompletion
+from sqlmesh.lsp.custom import (
+    AllModelsResponse,
+    MacroCompletion,
+    ModelCompletion,
+)
 from sqlmesh import macro
 import typing as t
 from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
-from sqlmesh.lsp.custom import ModelCompletion
 from sqlmesh.lsp.description import generate_markdown_description
 from sqlmesh.lsp.uri import URI
 
@@ -30,8 +33,10 @@ def get_sql_completions(
     # Combine keywords - SQL keywords first, then file keywords
     all_keywords = list(sql_keywords) + list(file_keywords - sql_keywords)
 
+    models = list(get_models(context, file_uri))
     return AllModelsResponse(
-        models=list(get_models(context, file_uri)),
+        models=[m.name for m in models],
+        model_completions=models,
         keywords=all_keywords,
         macros=list(get_macros(context, file_uri)),
     )

sqlmesh/lsp/custom.py

@@ -38,11 +38,11 @@ class ModelCompletion(PydanticModel):
 
 
 class AllModelsResponse(CustomMethodResponseBaseClass):
-    """
-    Response to get all the models that are in the current project.
-    """
+    """Response to get all models that are in the current project."""
 
-    models: t.List[ModelCompletion]
+    #: Deprecated: use ``model_completions`` instead
+    models: t.List[str]
+    model_completions: t.List[ModelCompletion]
     keywords: t.List[str]
     macros: t.List[MacroCompletion]
 

sqlmesh/lsp/main.py

@@ -680,7 +680,7 @@ def completion(
 
                 completion_items = []
                 # Add model completions
-                for model in completion_response.models:
+                for model in completion_response.model_completions:
                     completion_items.append(
                         types.CompletionItem(
                             label=model.name,

tests/lsp/test_completions.py

@@ -41,6 +41,20 @@ def test_get_macros():
     assert add_one_macro.description
 
 
+def test_model_completions_include_descriptions():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    completions = LSPContext.get_completions(lsp_context, None)
+
+    model_entry = next(
+        (m for m in completions.model_completions if m.name == "sushi.customers"),
+        None,
+    )
+    assert model_entry is not None
+    assert model_entry.description
+
+
 def test_get_sql_completions_with_context_no_file_uri():
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)

vscode/extension/src/completion/completion.ts

@@ -19,17 +19,10 @@ export const completionProvider = (
       if (isErr(result)) {
         return []
       }
-      const modelCompletions = result.value.models.map(model => {
-        const item = new vscode.CompletionItem(
-          model.name,
-          vscode.CompletionItemKind.Reference,
-        )
-        item.detail = 'SQLMesh Model'
-        if (model.description) {
-          item.documentation = new vscode.MarkdownString(model.description)
-        }
-        return item
-      })
+      const modelCompletions = result.value.models.map(
+        model =>
+          new vscode.CompletionItem(model, vscode.CompletionItemKind.Reference),
+      )
       const keywordCompletions = result.value.keywords.map(
         keyword =>
           new vscode.CompletionItem(keyword, vscode.CompletionItemKind.Keyword),

vscode/extension/src/lsp/custom.ts

@@ -25,16 +25,6 @@ export interface RenderModelEntry {
   rendered_query: string
 }
 
-export interface ModelCompletion {
-  name: string
-  description: string | null | undefined
-}
-
-export interface MacroCompletion {
-  name: string
-  description: string | null | undefined
-}
-
 // @eslint-disable-next-line  @typescript-eslint/consistent-type-definition
 export type CustomLSPMethods =
   | AllModelsMethod
@@ -51,9 +41,8 @@ interface AllModelsRequest {
 }
 
 interface AllModelsResponse {
-  models: ModelCompletion[]
+  models: string[]
   keywords: string[]
-  macros: MacroCompletion[]
 }
 
 export interface AbstractAPICallRequest {