Repository: TobikoData/sqlmesh
Pull request #4711 (2025-06-11 08:53:15)

feat: external model takes you to the model not just file


Patch:
diff --git a/examples/sushi/macros/utils.py b/examples/sushi/macros/utils.py
index a76bc3bf..fb2ccc21 100644
--- a/examples/sushi/macros/utils.py
+++ b/examples/sushi/macros/utils.py
@@ -5,14 +5,14 @@ from sqlmesh import SQL, macro
 
 @macro()
 def add_one(evaluator, column: int):
-    # typed column will be cast to an int and return an integer back
+    """typed column will be cast to an int and return an integer back"""
     assert isinstance(column, int)
     return column + 1
 
 
 @macro()
 def multiply(evaluator, column, num):
-    # untyped column will be a sqlglot column and return a sqlglot exp "column > 0"
+    """untyped column will be a sqlglot column and return a sqlglot exp "column > 0"""
     assert isinstance(column, exp.Column)
     return column * num
 
diff --git a/pyproject.toml b/pyproject.toml
index 6054adc0..b172c863 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,7 @@ dependencies = [
     "requests",
     "rich[jupyter]",
     "ruamel.yaml",
-    "sqlglot[rs]~=26.25.3",
+    "sqlglot[rs]~=26.26.0",
     "tenacity",
     "time-machine",
     "json-stream"
diff --git a/sqlmesh/lsp/completions.py b/sqlmesh/lsp/completions.py
index 1b1483c2..7e3781a5 100644
--- a/sqlmesh/lsp/completions.py
+++ b/sqlmesh/lsp/completions.py
@@ -1,6 +1,6 @@
 from functools import lru_cache
 from sqlglot import Dialect, Tokenizer
-from sqlmesh.lsp.custom import AllModelsResponse
+from sqlmesh.lsp.custom import AllModelsResponse, MacroCompletion
 from sqlmesh import macro
 import typing as t
 from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
@@ -60,15 +60,23 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.
     return all_models
 
 
-def get_macros(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
-    """Return a set of all macros with the ``@`` prefix."""
-    names = set(macro.get_registry())
+def get_macros(
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+) -> t.List[MacroCompletion]:
+    """Return a list of macros with optional descriptions."""
+    macros: t.Dict[str, t.Optional[str]] = {}
+
+    for name, m in macro.get_registry().items():
+        macros[name] = getattr(m.func, "__doc__", None)
+
     try:
         if context is not None:
-            names.update(context.context._macros)
+            for name, m in context.context._macros.items():
+                macros[name] = getattr(m.func, "__doc__", None)
     except Exception:
         pass
-    return names
+
+    return [MacroCompletion(name=name, description=doc) for name, doc in macros.items()]
 
 
 def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
diff --git a/sqlmesh/lsp/custom.py b/sqlmesh/lsp/custom.py
index c4328029..5c8123b7 100644
--- a/sqlmesh/lsp/custom.py
+++ b/sqlmesh/lsp/custom.py
@@ -23,6 +23,13 @@ class AllModelsRequest(CustomMethodRequestBaseClass):
     textDocument: types.TextDocumentIdentifier
 
 
+class MacroCompletion(PydanticModel):
+    """Information about a macro for autocompletion."""
+
+    name: str
+    description: t.Optional[str] = None
+
+
 class AllModelsResponse(CustomMethodResponseBaseClass):
     """
     Response to get all the models that are in the current project.
@@ -30,7 +37,7 @@ class AllModelsResponse(CustomMethodResponseBaseClass):
 
     models: t.List[str]
     keywords: t.List[str]
-    macros: t.List[str]
+    macros: t.List[MacroCompletion]
 
 
 RENDER_MODEL_FEATURE = "sqlmesh/render_model"
diff --git a/sqlmesh/lsp/helpers.py b/sqlmesh/lsp/helpers.py
index 7aa06ccb..920a93f5 100644
--- a/sqlmesh/lsp/helpers.py
+++ b/sqlmesh/lsp/helpers.py
@@ -6,22 +6,35 @@ from sqlmesh.core.linter.helpers import (
 )
 
 
-def to_lsp_range(
-    range: SQLMeshRange,
-) -> Range:
+def to_sqlmesh_position(position: Position) -> SQLMeshPosition:
     """
-    Converts a SQLMesh Range to an LSP Range.
+    Converts an LSP Position to a SQLMesh Position.
     """
-    return Range(
-        start=Position(line=range.start.line, character=range.start.character),
-        end=Position(line=range.end.line, character=range.end.character),
-    )
+    return SQLMeshPosition(line=position.line, character=position.character)
 
 
-def to_lsp_position(
-    position: SQLMeshPosition,
-) -> Position:
+def to_lsp_position(position: SQLMeshPosition) -> Position:
     """
     Converts a SQLMesh Position to an LSP Position.
     """
     return Position(line=position.line, character=position.character)
+
+
+def to_sqlmesh_range(range: Range) -> SQLMeshRange:
+    """
+    Converts an LSP Range to a SQLMesh Range.
+    """
+    return SQLMeshRange(
+        start=to_sqlmesh_position(range.start),
+        end=to_sqlmesh_position(range.end),
+    )
+
+
+def to_lsp_range(range: SQLMeshRange) -> Range:
+    """
+    Converts a SQLMesh Range to an LSP Range.
+    """
+    return Range(
+        start=to_lsp_position(range.start),
+        end=to_lsp_position(range.end),
+    )
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
index 3247bd5b..b028ea77 100755
--- a/sqlmesh/lsp/main.py
+++ b/sqlmesh/lsp/main.py
@@ -48,6 +48,7 @@ from sqlmesh.lsp.hints import get_hints
 from sqlmesh.lsp.reference import (
     LSPCteReference,
     LSPModelReference,
+    LSPExternalModelReference,
     get_references,
     get_all_references,
 )
@@ -444,11 +445,19 @@ class SQLMeshLanguageServer:
                 references = get_references(self.lsp_context, uri, params.position)
                 location_links = []
                 for reference in references:
-                    # Use target_range if available (CTEs, Macros), otherwise default to start of file
-                    if not isinstance(reference, LSPModelReference):
-                        target_range = reference.target_range
-                        target_selection_range = reference.target_range
-                    else:
+                    # Use target_range if available (CTEs, Macros, and external models in YAML)
+                    if isinstance(reference, LSPModelReference):
+                        # Regular SQL models - default to start of file
+                        target_range = types.Range(
+                            start=types.Position(line=0, character=0),
+                            end=types.Position(line=0, character=0),
+                        )
+                        target_selection_range = types.Range(
+                            start=types.Position(line=0, character=0),
+                            end=types.Position(line=0, character=0),
+                        )
+                    elif isinstance(reference, LSPExternalModelReference):
+                        # External models may have target_range set for YAML files
                         target_range = types.Range(
                             start=types.Position(line=0, character=0),
                             end=types.Position(line=0, character=0),
@@ -457,6 +466,13 @@ class SQLMeshLanguageServer:
                             start=types.Position(line=0, character=0),
                             end=types.Position(line=0, character=0),
                         )
+                        if reference.target_range is not None:
+                            target_range = reference.target_range
+                            target_selection_range = reference.target_range
+                    else:
+                        # CTEs and Macros always have target_range
+                        target_range = reference.target_range
+                        target_selection_range = reference.target_range
 
                     location_links.append(
                         types.LocationLink(
@@ -625,7 +641,8 @@ class SQLMeshLanguageServer:
                     and getattr(params.context, "trigger_character", None) == "@"
                 )
 
-                for macro_name in completion_response.macros:
+                for macro in completion_response.macros:
+                    macro_name = macro.name
                     insert_text = macro_name if triggered_by_at else f"@{macro_name}"
 
                     completion_items.append(
@@ -636,6 +653,7 @@ class SQLMeshLanguageServer:
                             filter_text=macro_name,
                             kind=types.CompletionItemKind.Function,
                             detail="SQLMesh Macro",
+                            documentation=macro.description,
                         )
                     )
 
diff --git a/sqlmesh/lsp/reference.py b/sqlmesh/lsp/reference.py
index e401be89..ac4d5374 100644
--- a/sqlmesh/lsp/reference.py
+++ b/sqlmesh/lsp/reference.py
@@ -8,7 +8,7 @@ from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.linter.helpers import (
     TokenPositionDetails,
 )
-from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.core.model.definition import SqlModel, ExternalModel
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
 from sqlglot import exp
 from sqlmesh.lsp.description import generate_markdown_description
@@ -22,10 +22,11 @@ import ast
 from sqlmesh.core.model import Model
 from sqlmesh import macro
 import inspect
+from ruamel.yaml import YAML
 
 
 class LSPModelReference(PydanticModel):
-    """A LSP reference to a model."""
+    """A LSP reference to a model, excluding external models."""
 
     type: t.Literal["model"] = "model"
     uri: str
@@ -33,6 +34,16 @@ class LSPModelReference(PydanticModel):
     markdown_description: t.Optional[str] = None
 
 
+class LSPExternalModelReference(PydanticModel):
+    """A LSP reference to an external model."""
+
+    type: t.Literal["external_model"] = "external_model"
+    uri: str
+    range: Range
+    markdown_description: t.Optional[str] = None
+    target_range: t.Optional[Range] = None
+
+
 class LSPCteReference(PydanticModel):
     """A LSP reference to a CTE."""
 
@@ -53,7 +64,8 @@ class LSPMacroReference(PydanticModel):
 
 
 Reference = t.Annotated[
-    t.Union[LSPModelReference, LSPCteReference, LSPMacroReference], Field(discriminator="type")
+    t.Union[LSPModelReference, LSPCteReference, LSPMacroReference, LSPExternalModelReference],
+    Field(discriminator="type"),
 ]
 
 
@@ -243,16 +255,38 @@ def get_model_definitions_for_a_path(
 
                 description = generate_markdown_description(referenced_model)
 
-                references.append(
-                    LSPModelReference(
-                        uri=referenced_model_uri.value,
-                        range=Range(
-                            start=to_lsp_position(start_pos_sqlmesh),
-                            end=to_lsp_position(end_pos_sqlmesh),
-                        ),
-                        markdown_description=description,
+                # For external models in YAML files, find the specific model block
+                if isinstance(referenced_model, ExternalModel):
+                    yaml_target_range: t.Optional[Range] = None
+                    if (
+                        referenced_model_path.suffix in (".yaml", ".yml")
+                        and referenced_model_path.is_file()
+                    ):
+                        yaml_target_range = _get_yaml_model_range(
+                            referenced_model_path, referenced_model.name
+                        )
+                    references.append(
+                        LSPExternalModelReference(
+                            uri=referenced_model_uri.value,
+                            range=Range(
+                                start=to_lsp_position(start_pos_sqlmesh),
+                                end=to_lsp_position(end_pos_sqlmesh),
+                            ),
+                            markdown_description=description,
+                            target_range=yaml_target_range,
+                        )
+                    )
+                else:
+                    references.append(
+                        LSPModelReference(
+                            uri=referenced_model_uri.value,
+                            range=Range(
+                                start=to_lsp_position(start_pos_sqlmesh),
+                                end=to_lsp_position(end_pos_sqlmesh),
+                            ),
+                            markdown_description=description,
+                        )
                     )
-                )
 
     return references
 
@@ -699,3 +733,31 @@ def _position_within_range(position: Position, range: Range) -> bool:
         range.end.line > position.line
         or (range.end.line == position.line and range.end.character >= position.character)
     )
+
+
+def _get_yaml_model_range(path: Path, model_name: str) -> t.Optional[Range]:
+    """
+    Find the range of a specific model block in a YAML file.
+
+    Args:
+        yaml_path: Path to the YAML file
+        model_name: Name of the model to find
+
+    Returns:
+        The Range of the model block in the YAML file, or None if not found
+    """
+    yaml = YAML()
+    with path.open("r", encoding="utf-8") as f:
+        data = yaml.load(f)
+
+    if not isinstance(data, list):
+        return None
+
+    for item in data:
+        if isinstance(item, dict) and item.get("name") == model_name:
+            # Get size of block by taking the earliest line/col in the items block and the last line/col of the block
+            position_data = item.lc.data["name"]  # type: ignore
+            start = Position(line=position_data[2], character=position_data[3])
+            end = Position(line=position_data[2], character=position_data[3] + len(item["name"]))
+            return Range(start=start, end=end)
+    return None


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index 76a5e66b..c6b23efa 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -643,9 +643,6 @@ def test_model_normalization_quote_flexibility():
         normalize_model_name('"catalog"."db"."table"', default_catalog=None, dialect="spark")
         == '"catalog"."db"."table"'
     )
-    # It doesn't work the other way which is what we currently expect
-    with pytest.raises(ParseError):
-        normalize_model_name("`catalog`.`db`.`table`", default_catalog=None, dialect=None)
 
 
 def test_macro_parse():
diff --git a/tests/lsp/test_completions.py b/tests/lsp/test_completions.py
index 8977d178..7e193d77 100644
--- a/tests/lsp/test_completions.py
+++ b/tests/lsp/test_completions.py
@@ -33,8 +33,12 @@ def test_get_macros():
     file_uri = URI.from_path(file_path)
     completions = LSPContext.get_completions(lsp_context, file_uri, file_content)
 
-    assert "each" in completions.macros
-    assert "add_one" in completions.macros
+    each_macro = next((m for m in completions.macros if m.name == "each"))
+    assert each_macro.name == "each"
+    assert each_macro.description
+    add_one_macro = next((m for m in completions.macros if m.name == "add_one"))
+    assert add_one_macro.name == "add_one"
+    assert add_one_macro.description
 
 
 def test_get_sql_completions_with_context_no_file_uri():
diff --git a/tests/lsp/test_reference_external_model.py b/tests/lsp/test_reference_external_model.py
new file mode 100644
index 00000000..ebf64209
--- /dev/null
+++ b/tests/lsp/test_reference_external_model.py
@@ -0,0 +1,38 @@
+from lsprotocol.types import Position
+from sqlmesh.core.context import Context
+from sqlmesh.core.linter.helpers import read_range_from_file
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.helpers import to_sqlmesh_range
+from sqlmesh.lsp.reference import get_references, LSPExternalModelReference
+from sqlmesh.lsp.uri import URI
+
+
+def test_reference() -> None:
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find model URIs
+    customers = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customers" in info.names
+    )
+
+    # Position of reference in file sushi.customers for sushi.raw_demographics
+    position = Position(line=42, character=20)
+    references = get_references(lsp_context, URI.from_path(customers), position)
+
+    assert len(references) == 1
+    reference = references[0]
+    assert isinstance(reference, LSPExternalModelReference)
+    assert reference.uri.endswith("external_models.yaml")
+
+    source_range = read_range_from_file(customers, to_sqlmesh_range(reference.range))
+    assert source_range == "raw.demographics"
+
+    if reference.target_range is None:
+        raise AssertionError("Reference target range should not be None")
+    target_range = read_range_from_file(
+        URI(reference.uri).to_path(), to_sqlmesh_range(reference.target_range)
+    )
+    assert target_range == "raw.demographics"