Feat(lsp): Add support for go to and find all References for CTEs (#4652)

Author: themisvaltinos

sqlmesh/lsp/main.py

@@ -36,9 +36,7 @@
     RenderModelRequest,
     RenderModelResponse,
 )
-from sqlmesh.lsp.reference import (
-    get_references,
-)
+from sqlmesh.lsp.reference import get_references, get_cte_references
 from sqlmesh.lsp.uri import URI
 from web.server.api.endpoints.lineage import column_lineage, model_lineage
 from web.server.api.endpoints.models import get_models
@@ -378,6 +376,28 @@ def goto_definition(
                 ls.show_message(f"Error getting references: {e}", types.MessageType.Error)
                 return []
 
+        @self.server.feature(types.TEXT_DOCUMENT_REFERENCES)
+        def find_references(
+            ls: LanguageServer, params: types.ReferenceParams
+        ) -> t.Optional[t.List[types.Location]]:
+            """Find all references of a symbol (currently supporting CTEs)"""
+            try:
+                uri = URI(params.text_document.uri)
+                self._ensure_context_for_document(uri)
+                document = ls.workspace.get_text_document(params.text_document.uri)
+                if self.lsp_context is None:
+                    raise RuntimeError(f"No context found for document: {document.path}")
+
+                cte_references = get_cte_references(self.lsp_context, uri, params.position)
+
+                # Convert references to Location objects
+                locations = [types.Location(uri=ref.uri, range=ref.range) for ref in cte_references]
+
+                return locations if locations else None
+            except Exception as e:
+                ls.show_message(f"Error getting locations: {e}", types.MessageType.Error)
+                return None
+
         @self.server.feature(types.TEXT_DOCUMENT_DIAGNOSTIC)
         def diagnostic(
             ls: LanguageServer, params: types.DocumentDiagnosticParams

sqlmesh/lsp/reference.py

@@ -47,16 +47,7 @@ def by_position(position: Position) -> t.Callable[[Reference], bool]:
     """
 
     def contains_position(r: Reference) -> bool:
-        return (
-            r.range.start.line < position.line
-            or (
-                r.range.start.line == position.line
-                and r.range.start.character <= position.character
-            )
-        ) and (
-            r.range.end.line > position.line
-            or (r.range.end.line == position.line and r.range.end.character >= position.character)
-        )
+        return _position_within_range(position, r.range)
 
     return contains_position
 
@@ -478,3 +469,78 @@ def get_built_in_macro_reference(macro_name: str, macro_range: Range) -> t.Optio
         ),
         markdown_description=func.__doc__ if func.__doc__ else None,
     )
+
+
+def get_cte_references(
+    lint_context: LSPContext, document_uri: URI, position: Position
+) -> t.List[Reference]:
+    """
+    Get all references to a CTE at a specific position in a document.
+
+    This function finds both the definition and all usages of a CTE within the same file.
+
+    Args:
+        lint_context: The LSP context
+        document_uri: The URI of the document
+        position: The position to check for CTE references
+
+    Returns:
+        A list of references to the CTE (including its definition and all usages)
+    """
+    references = get_model_definitions_for_a_path(lint_context, document_uri)
+
+    # Filter for CTE references (those with target_range set and same URI)
+    # TODO: Consider extending Reference class to explicitly indicate reference type instead
+    cte_references = [
+        ref for ref in references if ref.target_range is not None and ref.uri == document_uri.value
+    ]
+
+    if not cte_references:
+        return []
+
+    target_cte_definition_range = None
+    for ref in cte_references:
+        # Check if cursor is on a CTE usage
+        if _position_within_range(position, ref.range):
+            target_cte_definition_range = ref.target_range
+            break
+        # Check if cursor is on the CTE definition
+        elif ref.target_range and _position_within_range(position, ref.target_range):
+            target_cte_definition_range = ref.target_range
+            break
+
+    if target_cte_definition_range is None:
+        return []
+
+    # Add the CTE definition
+    matching_references = [
+        Reference(
+            uri=document_uri.value,
+            range=target_cte_definition_range,
+            markdown_description="CTE definition",
+        )
+    ]
+
+    # Add all usages
+    for ref in cte_references:
+        if ref.target_range == target_cte_definition_range:
+            matching_references.append(
+                Reference(
+                    uri=document_uri.value,
+                    range=ref.range,
+                    markdown_description="CTE usage",
+                )
+            )
+
+    return matching_references
+
+
+def _position_within_range(position: Position, range: Range) -> bool:
+    """Check if a position is within a given range."""
+    return (
+        range.start.line < position.line
+        or (range.start.line == position.line and range.start.character <= position.character)
+    ) and (
+        range.end.line > position.line
+        or (range.end.line == position.line and range.end.character >= position.character)
+    )

tests/lsp/test_reference_cte_find_all.py

@@ -0,0 +1,140 @@
+from lsprotocol.types import Position
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_cte_references
+from sqlmesh.lsp.uri import URI
+from tests.lsp.test_reference_cte import find_ranges_from_regex
+
+
+def test_cte_find_all_references():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    sushi_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customers" in info.names
+    )
+
+    with open(sushi_customers_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    # Test finding all references of "current_marketing"
+    ranges = find_ranges_from_regex(read_file, r"current_marketing(?!_outer)")
+    assert len(ranges) == 2
+
+    # Click on the CTE definition
+    position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4)
+    references = get_cte_references(lsp_context, URI.from_path(sushi_customers_path), position)
+
+    # Should find both the definition and the usage
+    assert len(references) == 2
+    assert all(ref.uri == URI.from_path(sushi_customers_path).value for ref in references)
+
+    reference_ranges = [ref.range for ref in references]
+    for expected_range in ranges:
+        assert any(
+            ref_range.start.line == expected_range.start.line
+            and ref_range.start.character == expected_range.start.character
+            for ref_range in reference_ranges
+        ), (
+            f"Expected to find reference at line {expected_range.start.line}, char {expected_range.start.character}"
+        )
+
+    # Click on the CTE usage
+    position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4)
+    references = get_cte_references(lsp_context, URI.from_path(sushi_customers_path), position)
+
+    # Should find the same references
+    assert len(references) == 2
+    assert all(ref.uri == URI.from_path(sushi_customers_path).value for ref in references)
+
+    reference_ranges = [ref.range for ref in references]
+    for expected_range in ranges:
+        assert any(
+            ref_range.start.line == expected_range.start.line
+            and ref_range.start.character == expected_range.start.character
+            for ref_range in reference_ranges
+        ), (
+            f"Expected to find reference at line {expected_range.start.line}, char {expected_range.start.character}"
+        )
+
+
+def test_cte_find_all_references_outer():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    sushi_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customers" in info.names
+    )
+
+    with open(sushi_customers_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    # Test finding all references of "current_marketing_outer"
+    ranges = find_ranges_from_regex(read_file, r"current_marketing_outer")
+    assert len(ranges) == 2
+
+    # Click on the CTE definition
+    position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4)
+    references = get_cte_references(lsp_context, URI.from_path(sushi_customers_path), position)
+
+    # Should find both the definition and the usage
+    assert len(references) == 2
+    assert all(ref.uri == URI.from_path(sushi_customers_path).value for ref in references)
+
+    # Verify that we found both occurrences
+    reference_ranges = [ref.range for ref in references]
+    for expected_range in ranges:
+        assert any(
+            ref_range.start.line == expected_range.start.line
+            and ref_range.start.character == expected_range.start.character
+            for ref_range in reference_ranges
+        ), (
+            f"Expected to find reference at line {expected_range.start.line}, char {expected_range.start.character}"
+        )
+
+    # Click on the CTE usage
+    position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4)
+    references = get_cte_references(lsp_context, URI.from_path(sushi_customers_path), position)
+
+    # Should find the same references
+    assert len(references) == 2
+    assert all(ref.uri == URI.from_path(sushi_customers_path).value for ref in references)
+
+    reference_ranges = [ref.range for ref in references]
+    for expected_range in ranges:
+        assert any(
+            ref_range.start.line == expected_range.start.line
+            and ref_range.start.character == expected_range.start.character
+            for ref_range in reference_ranges
+        ), (
+            f"Expected to find reference at line {expected_range.start.line}, char {expected_range.start.character}"
+        )
+
+
+def test_cte_no_references_on_non_cte():
+    # Test that clicking on non-CTE elements returns nothing, once this is supported adapt this test accordingly
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    sushi_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customers" in info.names
+    )
+
+    with open(sushi_customers_path, "r", encoding="utf-8") as file:
+        read_file = file.readlines()
+
+    # Click on a regular table reference
+    ranges = find_ranges_from_regex(read_file, r"sushi\.orders")
+    assert len(ranges) >= 1
+
+    position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4)
+    references = get_cte_references(lsp_context, URI.from_path(sushi_customers_path), position)
+
+    # Should find no references since this is not a CTE
+    assert len(references) == 0