feat(vscode): gtd for macros

Author: benfdking

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,19 +932,22 @@ 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):
-            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)  # type: ignore
-        return self.config_for_path(node._path)  # type: ignore
+            return self.config_for_path(self.get_snapshot(node, raise_if_missing=True).node._path)[
+                0
+            ]  # type: ignore
+        return self.config_for_path(node._path)[0]  # type: ignore
 
     @property
     def models(self) -> MappingProxyType[str, Model]:

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,10 @@
 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
+from sqlmesh import macro
+import inspect
 
 
 class Reference(PydanticModel):
@@ -72,6 +78,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
 
@@ -290,3 +301,180 @@ 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()
+        reference = get_macro_reference(
+            node=node,
+            target=target,
+            read_file=read_file,
+            config_path=config_path,
+            macro_name=macro_name,
+        )
+        if reference is not None:
+            references.append(reference)
+
+    return references
+
+
+def get_macro_reference(
+    target: t.Union[Model, StandaloneAudit],
+    read_file: t.List[str],
+    config_path: t.Optional[Path],
+    node: exp.Expression,
+    macro_name: str,
+) -> t.Optional[Reference]:
+    # Get the file path where the macro is defined
+    try:
+        # Get the position of the macro invocation in the source file first
+        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)
+
+            # Check if it's a built-in method
+            if builtin := get_built_in_macro_reference(macro_name, macro_range):
+                return builtin
+        else:
+            # Skip if we can't get the position
+            return None
+
+        # Find the macro definition information
+        macro_def = target.python_env.get(macro_name)
+        if macro_def is None:
+            return None
+
+        function_name = macro_def.name
+        if not function_name:
+            return None
+        if not macro_def.path:
+            return None
+        if not config_path:
+            return None
+        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())
+        with open(path, "r") as f:
+            output_read_line = f.readlines()
+
+        # Find the function definition by name
+        start_line = None
+        end_line = None
+        get_length_of_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
+                get_length_of_end_line = (
+                    len(output_read_line[end_line - 1])
+                    if end_line is not None and end_line - 1 < len(read_file)
+                    else 0
+                )
+                # Extract docstring if present
+                docstring = ast.get_docstring(ast_node)
+                break
+
+        if start_line is None or end_line is None or get_length_of_end_line is None:
+            return None
+
+        # Create a reference to the macro definition
+        macro_uri = URI.from_path(path)
+
+        return 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=get_length_of_end_line),
+            ),
+            markdown_description=docstring,
+        )
+    except Exception:
+        return None
+
+
+def get_built_in_macro_reference(macro_name: str, macro_range: Range) -> t.Optional[Reference]:
+    """
+    Get a reference to a built-in macro by its name.
+
+    Args:
+        macro_name: The name of the built-in macro (e.g., 'each', 'sql_literal')
+        macro_range: The range of the macro invocation in the source file
+    """
+    built_in_macros = macro.get_registry()
+    built_in_macro = built_in_macros.get(macro_name)
+    if built_in_macro is None:
+        return None
+
+    func = built_in_macro.func
+    filename = inspect.getfile(func)
+    source_lines, line_number = inspect.getsourcelines(func)
+
+    # Calculate the end line number by counting the number of source lines
+    end_line_number = line_number + len(source_lines) - 1
+
+    return Reference(
+        uri=URI.from_path(Path(filename)).value,
+        range=macro_range,
+        target_range=Range(
+            start=Position(line=line_number - 1, character=0),
+            end=Position(line=end_line_number - 1, character=0),
+        ),
+        markdown_description=func.__doc__ if func.__doc__ else None,
+    )

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("sushi/macros/utils.py")
+        assert ref.target_range is not None

tests/lsp/test_reference_macro_multi.py

@@ -0,0 +1,25 @@
+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
+    for ref in macro_references:
+        assert ref.uri.endswith("multi/repo_2/macros/__init__.py")
+        assert ref.target_range is not None

vscode/extension/tests/go_to_definition.spec.ts

@@ -0,0 +1,38 @@
+import { test, expect } from '@playwright/test';
+import path from 'path';
+import fs from 'fs-extra';
+import os from 'os';
+import { startVSCode, SUSHI_SOURCE_PATH } from './utils';
+
+test('Go to definition for macro', async () => {
+    const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));
+    await fs.copy(SUSHI_SOURCE_PATH, tempDir);
+
+    try {
+      const { window, close } = await startVSCode(tempDir);
+
+      //   Wait for the models folder to be visible
+      await window.waitForSelector('text=models');
+
+      // Click on the models folder, excluding external_models
+      await window.getByRole('treeitem', { name: 'models', exact: true }).locator('a').click();
+
+      // Open the customer_revenue_lifetime model
+      await window.getByRole('treeitem', { name: 'top_waiters.sql', exact: true }).locator('a').click();
+
+      await window.waitForSelector('text=grain');
+      await window.waitForSelector('text=Loaded SQLMesh Context')
+
+      // Render the model
+      window.locator("text=@MULTIPLY").click({
+        modifiers: ["Meta"]
+      })
+
+      // Check if the model is rendered by check if "`oi`.`order_id` AS `order_id`," is in the window
+      await expect(window.locator('text=def multiply(')).toBeVisible();
+
+      await close();
+    } finally {
+      await fs.removeSync(tempDir);
+    }
+  });

vscode/extension/tests/utils.ts

@@ -40,11 +40,8 @@ export const startVSCode = async (workspaceDir: string): Promise<{
       args,
     });
     const window = await electronApp.firstWindow();
-    await window.waitForLoadState('domcontentloaded');
-    await window.waitForLoadState('networkidle');
-    await window.waitForTimeout(2_000);
     return { window, close: async () => {
       await electronApp.close();
-      await fs.remove(userDataDir);
+      await fs.removeSync(userDataDir);
     } };
   }

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