Feat(lsp): Add find all and go to references support for Macros (#4692)

Author: themisvaltinos

examples/sushi/models/customers.sql

@@ -33,7 +33,8 @@ LEFT JOIN (
   WITH current_marketing AS (
     SELECT
       customer_id,
-      status
+      status,
+      @ADD_ONE(1) AS another_column,
     FROM current_marketing_outer
   )
   SELECT * FROM current_marketing

sqlmesh/lsp/reference.py

@@ -450,75 +450,76 @@ def get_model_find_all_references(
     Returns:
         A list of references to the model across all files
     """
-    # First, get the references in the current file to determine what model we're looking for
-    current_file_references = [
-        ref
-        for ref in get_model_definitions_for_a_path(lint_context, document_uri)
-        if isinstance(ref, LSPModelReference)
-    ]
-
     # Find the model reference at the cursor position
-    target_model_uri: t.Optional[str] = None
-    for ref in current_file_references:
-        if _position_within_range(position, ref.range):
-            # This is a model reference, get the target model URI
-            target_model_uri = ref.uri
-            break
+    model_at_position = next(
+        filter(
+            lambda ref: isinstance(ref, LSPModelReference)
+            and _position_within_range(position, ref.range),
+            get_model_definitions_for_a_path(lint_context, document_uri),
+        ),
+        None,
+    )
 
-    if target_model_uri is None:
+    if not model_at_position:
         return []
 
+    assert isinstance(model_at_position, LSPModelReference)  # for mypy
+
+    target_model_uri = model_at_position.uri
+
     # Start with the model definition
     all_references: t.List[LSPModelReference] = [
         LSPModelReference(
-            uri=ref.uri,
+            uri=model_at_position.uri,
             range=Range(
                 start=Position(line=0, character=0),
                 end=Position(line=0, character=0),
             ),
-            markdown_description=ref.markdown_description,
+            markdown_description=model_at_position.markdown_description,
         )
     ]
 
-    # Then add the original reference
-    for ref in current_file_references:
-        if ref.uri == target_model_uri and isinstance(ref, LSPModelReference):
-            all_references.append(
-                LSPModelReference(
-                    uri=document_uri.value,
-                    range=ref.range,
-                    markdown_description=ref.markdown_description,
-                )
+    # Then add references from the current file
+    current_file_refs = filter(
+        lambda ref: isinstance(ref, LSPModelReference) and ref.uri == target_model_uri,
+        get_model_definitions_for_a_path(lint_context, document_uri),
+    )
+
+    for ref in current_file_refs:
+        assert isinstance(ref, LSPModelReference)  # for mypy
+
+        all_references.append(
+            LSPModelReference(
+                uri=document_uri.value,
+                range=ref.range,
+                markdown_description=ref.markdown_description,
             )
+        )
 
     # Search through the models in the project
-    for path, target in lint_context.map.items():
-        if not isinstance(target, (ModelTarget, AuditTarget)):
-            continue
-
+    for path, _ in lint_context.map.items():
         file_uri = URI.from_path(path)
 
         # Skip current file, already processed
         if file_uri.value == document_uri.value:
             continue
 
-        # Get model references for this file
-        file_references = [
-            ref
-            for ref in get_model_definitions_for_a_path(lint_context, file_uri)
-            if isinstance(ref, LSPModelReference)
-        ]
-
-        # Add references that point to the target model file
-        for ref in file_references:
-            if ref.uri == target_model_uri and isinstance(ref, LSPModelReference):
-                all_references.append(
-                    LSPModelReference(
-                        uri=file_uri.value,
-                        range=ref.range,
-                        markdown_description=ref.markdown_description,
-                    )
+        # Get model references that point to the target model
+        matching_refs = filter(
+            lambda ref: isinstance(ref, LSPModelReference) and ref.uri == target_model_uri,
+            get_model_definitions_for_a_path(lint_context, file_uri),
+        )
+
+        for ref in matching_refs:
+            assert isinstance(ref, LSPModelReference)  # for mypy
+
+            all_references.append(
+                LSPModelReference(
+                    uri=file_uri.value,
+                    range=ref.range,
+                    markdown_description=ref.markdown_description,
                 )
+            )
 
     return all_references
 
@@ -587,13 +588,83 @@ def get_cte_references(
     return matching_references
 
 
+def get_macro_find_all_references(
+    lsp_context: LSPContext, document_uri: URI, position: Position
+) -> t.List[LSPMacroReference]:
+    """
+    Get all references to a macro at a specific position in a document.
+
+    This function finds all usages of a macro across the entire project.
+
+    Args:
+        lsp_context: The LSP context
+        document_uri: The URI of the document
+        position: The position to check for macro references
+
+    Returns:
+        A list of references to the macro across all files
+    """
+    # Find the macro reference at the cursor position
+    macro_at_position = next(
+        filter(
+            lambda ref: isinstance(ref, LSPMacroReference)
+            and _position_within_range(position, ref.range),
+            get_macro_definitions_for_a_path(lsp_context, document_uri),
+        ),
+        None,
+    )
+
+    if not macro_at_position:
+        return []
+
+    assert isinstance(macro_at_position, LSPMacroReference)  # for mypy
+
+    target_macro_uri = macro_at_position.uri
+    target_macro_target_range = macro_at_position.target_range
+
+    # Start with the macro definition
+    all_references: t.List[LSPMacroReference] = [
+        LSPMacroReference(
+            uri=target_macro_uri,
+            range=target_macro_target_range,
+            target_range=target_macro_target_range,
+            markdown_description=None,
+        )
+    ]
+
+    # Search through all SQL and audit files in the project
+    for path, _ in lsp_context.map.items():
+        file_uri = URI.from_path(path)
+
+        # Get macro references that point to the same macro definition
+        matching_refs = filter(
+            lambda ref: isinstance(ref, LSPMacroReference)
+            and ref.uri == target_macro_uri
+            and ref.target_range == target_macro_target_range,
+            get_macro_definitions_for_a_path(lsp_context, file_uri),
+        )
+
+        for ref in matching_refs:
+            assert isinstance(ref, LSPMacroReference)  # for mypy
+            all_references.append(
+                LSPMacroReference(
+                    uri=file_uri.value,
+                    range=ref.range,
+                    target_range=ref.target_range,
+                    markdown_description=ref.markdown_description,
+                )
+            )
+
+    return all_references
+
+
 def get_all_references(
     lint_context: LSPContext, document_uri: URI, position: Position
 ) -> t.Sequence[Reference]:
     """
     Get all references of a symbol at a specific position in a document.
 
-    This function determines the type of reference (CTE, model for now) at the cursor
+    This function determines the type of reference (CTE, model or macro) at the cursor
     position and returns all references to that symbol across the project.
 
     Args:
@@ -612,6 +683,10 @@ def get_all_references(
     if model_references := get_model_find_all_references(lint_context, document_uri, position):
         return model_references
 
+    # Finally try macro references (across files)
+    if macro_references := get_macro_find_all_references(lint_context, document_uri, position):
+        return macro_references
+
     return []
 
 

tests/integrations/github/cicd/test_integration.py

@@ -1684,7 +1684,7 @@ def test_overlapping_changes_models(
 
 +++ 
 
-@@ -25,7 +25,8 @@
+@@ -29,7 +29,8 @@
 
  SELECT DISTINCT
    CAST(o.customer_id AS INT) AS customer_id,