Fix macro completions

Author: benfdking

sqlmesh/lsp/completions.py

@@ -59,8 +59,14 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.
 
 
 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())
+    """Return a set of all macros with the ``@`` prefix."""
+    names = set(macro.get_registry().keys())
+    if context is not None:
+        try:
+            names.update(context.context._macros.keys())
+        except Exception:
+            pass
+    return {f"@{name}" for name in names}
 
 
 def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:

sqlmesh/lsp/main.py

@@ -521,6 +521,15 @@ def completion(
                             detail="SQLMesh Model",
                         )
                     )
+                # Add macro completions
+                for macro_name in completion_response.macros:
+                    completion_items.append(
+                        types.CompletionItem(
+                            label=macro_name,
+                            kind=types.CompletionItemKind.Function,
+                            detail="SQLMesh Macro",
+                        )
+                    )
                 # Add keyword completions
                 for keyword in completion_response.keywords:
                     completion_items.append(

tests/lsp/test_completions.py

@@ -20,7 +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
+    assert "@each" in completions.macros
 
 
 def test_get_sql_completions_with_context_no_file_uri():
@@ -31,7 +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
+    assert "@add_one" in completions.macros
 
 
 def test_get_sql_completions_with_context_and_file_uri():
@@ -42,7 +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
+    assert "@add_one" in completions.macros
 
 
 def test_extract_keywords_from_content():