feat(vscode): gtd for macros

Author: benfdking

examples/multi/repo_1/config.yaml

@@ -10,7 +10,7 @@ gateways:
     connection:
       type: duckdb
 
-default_gateway: local
+default_gateway: memory
 
 
 before_all:

examples/multi/repo_2/config.yaml

@@ -10,7 +10,7 @@ gateways:
     connection:
       type: duckdb
 
-default_gateway: local
+default_gateway: memory
 
 
 before_all:

examples/sushi/macros/utils.py

@@ -26,6 +26,7 @@ def sql_literal(
     column_str: str,
     column_quoted: str,
 ):
+    """A macro that accepts various types of SQL literals and returns the column."""
     assert isinstance(column, str)
     assert isinstance(str_lit, str)
     assert str_lit == "'x'"

sqlmesh/core/context.py

@@ -932,14 +932,15 @@ def get_snapshot(
 
         return snapshot
 
-    def config_for_path(self, path: Path) -> Config:
+    def config_for_path(self, path: Path) -> t.Tuple[Config, Path]:
+        """Returns the config and path of the said project for a given file path."""
         for config_path, config in self.configs.items():
             try:
                 path.relative_to(config_path)
-                return config
+                return config, config_path
             except ValueError:
                 pass
-        return self.config
+        return self.config, self.path
 
     def config_for_node(self, node: str | Model | Audit) -> Config:
         if isinstance(node, str):

sqlmesh/lsp/reference.py

@@ -1,6 +1,8 @@
 from lsprotocol.types import Range, Position
 import typing as t
+from pathlib import Path
 
+from sqlmesh.core.audit import StandaloneAudit
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
@@ -10,6 +12,8 @@
 from sqlmesh.lsp.uri import URI
 from sqlmesh.utils.pydantic import PydanticModel
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+import ast
+from sqlmesh.core.model import Model
 
 
 class Reference(PydanticModel):
@@ -72,6 +76,11 @@ def get_references(
         A list of references at the given position
     """
     references = get_model_definitions_for_a_path(lint_context, document_uri)
+
+    # Get macro references before filtering by position
+    macro_references = get_macro_definitions_for_a_path(lint_context, document_uri)
+    references.extend(macro_references)
+
     filtered_references = list(filter(by_position(position), references))
     return filtered_references
 
@@ -287,3 +296,118 @@ def _range_from_token_position_details(
         start=Position(line=start_line_0, character=start_col_0),
         end=Position(line=end_line_0, character=end_col_0),
     )
+
+
+def get_macro_definitions_for_a_path(
+    lsp_context: LSPContext, document_uri: URI
+) -> t.List[Reference]:
+    """
+    Get macro references for a given path.
+
+    This function finds all macro invocations (e.g., @ADD_ONE, @MULTIPLY) in a SQL file
+    and creates references to their definitions in the Python macro files.
+
+    Args:
+        lsp_context: The LSP context containing macro definitions
+        document_uri: The URI of the document to search for macro invocations
+
+    Returns:
+        A list of Reference objects for each macro invocation found
+    """
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
+        return []
+
+    # Get the file info from the context map
+    if path not in lsp_context.map:
+        return []
+
+    file_info = lsp_context.map[path]
+    # Process based on whether it's a model or standalone audit
+    if isinstance(file_info, ModelTarget):
+        # It's a model
+        target: t.Optional[t.Union[Model, StandaloneAudit]] = lsp_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if target is None or not isinstance(target, SqlModel):
+            return []
+        query = target.query
+        file_path = target._path
+    elif isinstance(file_info, AuditTarget):
+        # It's a standalone audit
+        target = lsp_context.context.standalone_audits.get(file_info.name)
+        if target is None:
+            return []
+        query = target.query
+        file_path = target._path
+    else:
+        return []
+
+    references = []
+    config_for_model, config_path = lsp_context.context.config_for_path(
+        file_path,
+    )
+
+    with open(file_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    for node in query.find_all(exp.Anonymous):
+        macro_name = node.name.lower()
+
+        # Find the macro definition information
+        macro_def = target.python_env.get(macro_name)
+        if macro_def is None:
+            continue
+
+        # Get the file path where the macro is defined
+        try:
+            function_name = macro_def.name
+            if not function_name:
+                continue
+            if not macro_def.path:
+                continue
+            path = Path(config_path).joinpath(macro_def.path)
+
+            # Parse the Python file to find the function definition
+            with open(path, "r") as f:
+                tree = ast.parse(f.read())
+
+            # Find the function definition by name
+            start_line = None
+            end_line = None
+            docstring = None
+            for ast_node in ast.walk(tree):
+                if isinstance(ast_node, ast.FunctionDef) and ast_node.name == function_name:
+                    start_line = ast_node.lineno
+                    end_line = ast_node.end_lineno
+                    # Extract docstring if present
+                    docstring = ast.get_docstring(ast_node)
+                    break
+
+            if start_line is None or end_line is None:
+                continue
+
+            # Create a reference to the macro definition
+            macro_uri = URI.from_path(path)
+
+            # Get the position of the macro invocation in the source file
+            if hasattr(node, "meta") and node.meta:
+                token_details = TokenPositionDetails.from_meta(node.meta)
+                macro_range = _range_from_token_position_details(token_details, read_file)
+
+                references.append(
+                    Reference(
+                        uri=macro_uri.value,
+                        range=macro_range,
+                        target_range=Range(
+                            start=Position(line=start_line - 1, character=0),
+                            end=Position(line=end_line - 1, character=0),
+                        ),
+                        markdown_description=docstring,
+                    )
+                )
+        except (OSError, TypeError):
+            # If we can't get the source file, skip this macro
+            continue
+
+    return references

tests/lsp/test_reference_macro.py

@@ -0,0 +1,30 @@
+import pytest
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_macro_definitions_for_a_path
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_macro_references() -> None:
+    """Test that macro references (e.g., @ADD_ONE, @MULTIPLY) have proper go-to-definition support."""
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find the top_waiters model that uses macros
+    top_waiters_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.top_waiters" in info.names
+    )
+
+    top_waiters_uri = URI.from_path(top_waiters_path)
+    macro_references = get_macro_definitions_for_a_path(lsp_context, top_waiters_uri)
+
+    # We expect 3 macro references: @ADD_ONE, @MULTIPLY, @SQL_LITERAL
+    assert len(macro_references) == 3
+
+    # Check that all references point to the utils.py file
+    for ref in macro_references:
+        assert ref.uri.endswith("macros/utils.py")
+        assert ref.target_range is not None  # Should have target range for go-to-definition

tests/lsp/test_reference_macro_multi.py

@@ -0,0 +1,22 @@
+import pytest
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_macro_definitions_for_a_path
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_macro_references_multirepo() -> None:
+    context = Context(paths=["examples/multi/repo_1", "examples/multi/repo_2"])
+    lsp_context = LSPContext(context)
+
+    d_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "silver.d" in info.names
+    )
+
+    d = URI.from_path(d_path)
+    macro_references = get_macro_definitions_for_a_path(lsp_context, d)
+
+    assert len(macro_references) == 2

web/server/api/endpoints/files.py

@@ -60,7 +60,7 @@ async def write_file(
         replace_file(settings.project_path / path, settings.project_path / path_or_new_path)
     else:
         full_path = settings.project_path / path
-        config = context.config_for_path(Path(path_or_new_path)) if context else None
+        config, _ = context.config_for_path(Path(path_or_new_path)) if context else (None, None)
         if (
             config
             and config.ui.format_on_save

web/server/utils.py

@@ -64,11 +64,10 @@ def validate_path(path: str, settings: Settings = Depends(get_settings)) -> str:
     if any(
         full_path.match(pattern)
         for pattern in (
-            context.config_for_path(Path(path)).ignore_patterns if context else c.IGNORE_PATTERNS
+            context.config_for_path(Path(path))[0].ignore_patterns if context else c.IGNORE_PATTERNS
         )
     ):
         raise HTTPException(status_code=HTTP_404_NOT_FOUND)
-
     return path