temp

Author: benfdking

sqlmesh/lsp/main.py

@@ -280,7 +280,7 @@ def goto_definition(
                             start=types.Position(line=0, character=0),
                             end=types.Position(line=0, character=0),
                         )
-                    
+
                     location_links.append(
                         types.LocationLink(
                             target_uri=reference.uri,

sqlmesh/lsp/reference.py

@@ -130,7 +130,7 @@ def get_model_definitions_for_a_path(
 
     # Get SQL query and find all table references
     tables = list(query.find_all(exp.Table))
-    
+
     with open(file_path, "r", encoding="utf-8") as file:
         read_file = file.readlines()
 
@@ -139,33 +139,33 @@ def get_model_definitions_for_a_path(
     with_clause = query.find(exp.With)
     if with_clause:
         for cte in with_clause.expressions:
-            if cte.alias and hasattr(cte, 'meta') and cte.meta:
+            if cte.alias and hasattr(cte, "meta") and cte.meta:
                 cte_definitions[cte.alias] = cte
 
     for table in tables:
         table_name = table.name
-        
+
         # Check if this table reference is a CTE
         if table_name in cte_definitions:
             # This is a CTE reference - create a reference to the CTE definition
             cte_def = cte_definitions[table_name]
-            if hasattr(cte_def, 'meta') and cte_def.meta:
+            if hasattr(cte_def, "meta") and cte_def.meta:
                 try:
                     # Get the position of the table reference
                     table_meta = TokenPositionDetails.from_meta(table.this.meta)
                     table_range = _range_from_token_position_details(table_meta, read_file)
-                    
+
                     # Get the position of the CTE definition (alias part)
                     cte_meta = TokenPositionDetails.from_meta(cte_def.meta)
                     cte_range = _range_from_token_position_details(cte_meta, read_file)
-                    
+
                     # Create a reference from the table usage to the CTE definition
                     references.append(
                         Reference(
                             uri=document_uri.value,  # Same file
                             range=table_range,
                             description=f"CTE: {table_name}",
-                            target_range=cte_range
+                            target_range=cte_range,
                         )
                     )
                 except Exception:

tests/lsp/simple_cte_test.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Test script for CTE go-to-definition functionality."""
+
+import re
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.reference import get_references
+from sqlmesh.lsp.uri import URI
+from lsprotocol.types import Range, Position
+
+def test_cte_parsing():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find model URIs
+    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()
+
+    references = get_references(lsp_context, URI.from_path(sushi_customers_path), Position(line=0, character=0))
+
+    assert len(references) == 1
+    assert references[0].uri == URI.from_path(sushi_customers_path)
+
+    ranges = find_ranges_from_regex(read_file, r"WITH\s+current_marketing\s+AS\s+\(SELECT\s+customer_id,\s+status\s+FROM\s+sushi\.marketing\s+WHERE\s+valid_to\s+is\s+null\)")
+    assert len(ranges) == 1
+    assert ranges[0].start == Position(line=1, character=0)
+    assert ranges[0].end == Position(line=1, character=1)
+
+
+
+
+
+def find_ranges_from_regex(read_file: t.List[str], regex: str) -> t.List[Range]:
+    """Find all ranges in the read file that match the regex."""
+    return [
+        Range(
+            start=Position(line=line_number, character=match.start()),
+            end=Position(line=line_number, character=match.end())
+        )
+        for line_number, line in enumerate(read_file)
+        for match in [m for m in [re.search(regex, line)] if m]
+    ]
+
+