feat(lsp): include macros in autocomplete

Author: benfdking

sqlmesh/lsp/completions.py

@@ -1,6 +1,7 @@
 from functools import lru_cache
 from sqlglot import Dialect, Tokenizer
 from sqlmesh.lsp.custom import AllModelsResponse
+from sqlmesh import macro
 import typing as t
 from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
 from sqlmesh.lsp.uri import URI
@@ -26,6 +27,7 @@ def get_sql_completions(
     return AllModelsResponse(
         models=list(get_models(context, file_uri)),
         keywords=all_keywords,
+        macros=list(get_macros(context)),
     )
 
 
@@ -56,6 +58,11 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.
     return all_models
 
 
+def get_macros(context: t.Optional[LSPContext]) -> t.Set[str]:
+    """Return a list of all macros available in the context."""
+    return set(macro.get_registry().keys())
+
+
 def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
     """
     Return a list of sql keywords for a given file.
@@ -159,3 +166,4 @@ def extract_keywords_from_content(content: str, dialect: t.Optional[str] = None)
         pass
 
     return keywords
+

sqlmesh/lsp/custom.py

@@ -20,6 +20,7 @@ class AllModelsResponse(PydanticModel):
 
     models: t.List[str]
     keywords: t.List[str]
+    macros: t.List[str]
 
 
 RENDER_MODEL_FEATURE = "sqlmesh/render_model"

tests/lsp/test_completions.py

@@ -20,6 +20,7 @@ def test_get_sql_completions_no_context():
     completions = get_sql_completions(None, None)
     assert len(completions.keywords) >= len(TOKENIZER_KEYWORDS)
     assert len(completions.models) == 0
+    assert "each" in completions.macros
 
 
 def test_get_sql_completions_with_context_no_file_uri():
@@ -30,6 +31,7 @@ def test_get_sql_completions_with_context_no_file_uri():
     assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
     assert "sushi.active_customers" in completions.models
     assert "sushi.customers" in completions.models
+    assert "add_one" in completions.macros
 
 
 def test_get_sql_completions_with_context_and_file_uri():
@@ -40,6 +42,7 @@ def test_get_sql_completions_with_context_and_file_uri():
     completions = lsp_context.get_autocomplete(URI.from_path(file_uri))
     assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
     assert "sushi.active_customers" not in completions.models
+    assert "add_one" in completions.macros
 
 
 def test_extract_keywords_from_content():

vscode/extension/src/completion/completion.ts

@@ -23,12 +23,17 @@ export const completionProvider = (
         model =>
           new vscode.CompletionItem(model, vscode.CompletionItemKind.Reference),
       )
+      const macroCompletions = result.value.macros.map(
+        macro =>
+          new vscode.CompletionItem(macro, vscode.CompletionItemKind.Function),
+      )
       const keywordCompletions = result.value.keywords.map(
         keyword =>
           new vscode.CompletionItem(keyword, vscode.CompletionItemKind.Keyword),
       )
       return new vscode.CompletionList([
         ...modelCompletions,
+        ...macroCompletions,
         ...keywordCompletions,
       ])
     },

vscode/extension/src/lsp/custom.ts

@@ -41,6 +41,7 @@ interface AllModelsRequest {
 interface AllModelsResponse {
   models: string[]
   keywords: string[]
+  macros: string[]
 }
 
 export interface AbstractAPICallRequest {