feat(vscode): gtd for macros

[ci skip]

Author: benfdking

sqlmesh/lsp/reference.py

@@ -1,5 +1,6 @@
 from lsprotocol.types import Range, Position
 import typing as t
+from pathlib import Path
 
 from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.model.definition import SqlModel
@@ -72,6 +73,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 +293,100 @@ 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(
+    lint_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:
+        lint_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 lint_context.map:
+        return []
+    
+    file_info = lint_context.map[path]
+    # Process based on whether it's a model or standalone audit
+    if isinstance(file_info, ModelTarget):
+        # It's a model
+        model = lint_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if model is None or not isinstance(model, SqlModel):
+            return []
+        query = model.query
+        file_path = model._path
+    elif isinstance(file_info, AuditTarget):
+        # It's a standalone audit
+        audit = lint_context.context.standalone_audits.get(file_info.name)
+        if audit is None:
+            return []
+        query = audit.query
+        file_path = audit._path
+    else:
+        return []
+    
+    references = []
+    
+    with open(file_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+    
+    # Find all macro function calls in the query
+    from sqlmesh.core.macros import MacroFunc
+    
+    for node in query.find_all(MacroFunc):
+        macro_name = node.name.lower()
+        
+        # Find the macro definition in the context
+        macro_def = lint_context.context._macros.get(macro_name)
+        if macro_def is None:
+            continue
+            
+        # Get the file path where the macro is defined
+        import inspect
+        try:
+            source_file = inspect.getsourcefile(macro_def)
+            if source_file is None:
+                continue
+                
+            # Get the line number where the macro is defined
+            source_lines, start_line = inspect.getsourcelines(macro_def)
+            
+            # Create a reference to the macro definition
+            macro_uri = URI.from_path(Path(source_file))
+            
+            # 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=start_line - 1, character=0)
+                        ),
+                        markdown_description=f"Macro: `@{macro_name.upper()}`"
+                    )
+                )
+        except (OSError, TypeError):
+            # If we can't get the source file, skip this macro
+            continue
+    
+    return references

tests/lsp/test_macro_reference.py

@@ -0,0 +1,124 @@
+import pytest
+from lsprotocol.types import Position
+from pathlib import Path
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_macro_definitions_for_a_path, get_references, by_position
+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
+    
+    # Read the SQL file to verify macro positions
+    with open(top_waiters_path, "r") as file:
+        content = file.read()
+    
+    # Verify that the references are at the expected macro invocations
+    assert "@ADD_ONE" in content
+    assert "@MULTIPLY" in content
+    assert "@SQL_LITERAL" in content
+
+
+@pytest.mark.fast 
+def test_macro_go_to_definition_with_position() -> None:
+    """Test go-to-definition for specific macro positions."""
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+    
+    # Find the top_waiters model
+    top_waiters_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.top_waiters" in info.names
+    )
+    
+    # Read the file to find exact positions
+    with open(top_waiters_path, "r") as file:
+        lines = file.readlines()
+    
+    # Find line with @ADD_ONE
+    add_one_line = None
+    add_one_char = None
+    for i, line in enumerate(lines):
+        if "@ADD_ONE" in line:
+            add_one_line = i
+            add_one_char = line.index("@ADD_ONE") + 4  # Position inside the macro name
+            break
+    
+    assert add_one_line is not None
+    
+    # Test go-to-definition at the @ADD_ONE position
+    position = Position(line=add_one_line, character=add_one_char)
+    references = get_references(lsp_context, URI.from_path(top_waiters_path), position)
+    
+    # Should find exactly one reference
+    assert len(references) == 1
+    assert references[0].uri.endswith("macros/utils.py")
+    assert references[0].target_range is not None
+    
+    # The target should be at the line where add_one is defined
+    # (approximately line 6 based on the utils.py file)
+    assert references[0].target_range.start.line >= 5  # Zero-indexed, so line 6 is index 5
+
+
+@pytest.mark.fast
+def test_all_three_macros_positions() -> None:
+    """Test that all three macros (@ADD_ONE, @MULTIPLY, @SQL_LITERAL) work with position-based lookup."""
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+    
+    # Find the top_waiters model
+    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)
+    
+    # Get all references (models and macros)
+    with open(top_waiters_path, "r") as file:
+        lines = file.readlines()
+    
+    macro_names = ["@ADD_ONE", "@MULTIPLY", "@SQL_LITERAL"]
+    found_macros = []
+    
+    for macro_name in macro_names:
+        for i, line in enumerate(lines):
+            if macro_name in line:
+                char_pos = line.index(macro_name) + len(macro_name) // 2
+                position = Position(line=i, character=char_pos)
+                refs = get_references(lsp_context, top_waiters_uri, position)
+                
+                # Filter to only macro references
+                macro_refs = [r for r in refs if r.uri.endswith("utils.py")]
+                if macro_refs:
+                    found_macros.append(macro_name)
+                    assert len(macro_refs) == 1
+                    assert macro_refs[0].markdown_description == f"Macro: `{macro_name}`"
+    
+    # Should find all three macros
+    assert len(found_macros) == 3
+    assert set(found_macros) == set(macro_names)