Skip to content

Commit 51481b2

Browse files
authored
feat(lsp): references function now detects non registered external models (#5023)
1 parent c02022e commit 51481b2

File tree

4 files changed

+65
-14
lines changed

4 files changed

+65
-14
lines changed

sqlmesh/lsp/main.py

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -496,14 +496,15 @@ def goto_definition(
496496
target_range = reference.target_range
497497
target_selection_range = reference.target_range
498498

499-
location_links.append(
500-
types.LocationLink(
501-
target_uri=reference.uri,
502-
target_selection_range=target_selection_range,
503-
target_range=target_range,
504-
origin_selection_range=reference.range,
499+
if reference.uri is not None:
500+
location_links.append(
501+
types.LocationLink(
502+
target_uri=reference.uri,
503+
target_selection_range=target_selection_range,
504+
target_range=target_range,
505+
origin_selection_range=reference.range,
506+
)
505507
)
506-
)
507508
return location_links
508509
except Exception as e:
509510
ls.show_message(f"Error getting references: {e}", types.MessageType.Error)
@@ -521,7 +522,11 @@ def find_references(
521522
all_references = get_all_references(context, uri, params.position)
522523

523524
# Convert references to Location objects
524-
locations = [types.Location(uri=ref.uri, range=ref.range) for ref in all_references]
525+
locations = [
526+
types.Location(uri=ref.uri, range=ref.range)
527+
for ref in all_references
528+
if ref.uri is not None
529+
]
525530

526531
return locations if locations else None
527532
except Exception as e:

sqlmesh/lsp/reference.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,13 @@ class LSPExternalModelReference(PydanticModel):
3838
"""A LSP reference to an external model."""
3939

4040
type: t.Literal["external_model"] = "external_model"
41-
uri: str
4241
range: Range
43-
markdown_description: t.Optional[str] = None
4442
target_range: t.Optional[Range] = None
43+
uri: t.Optional[str] = None
44+
"""The URI of the external model, typically a YAML file, it is optional because
45+
external models can be unregistered and so they URI is not available."""
46+
47+
markdown_description: t.Optional[str] = None
4548

4649

4750
class LSPCteReference(PydanticModel):
@@ -224,7 +227,7 @@ def get_model_definitions_for_a_path(
224227
references.extend(column_references)
225228
continue
226229

227-
# For non-CTE tables, process as before (external model references)
230+
# For non-CTE tables, process these as before (external model references)
228231
# Normalize the table reference
229232
unaliased = table.copy()
230233
if unaliased.args.get("alias") is not None:
@@ -247,6 +250,19 @@ def get_model_definitions_for_a_path(
247250
model_or_snapshot=normalized_reference_name, raise_if_missing=False
248251
)
249252
if referenced_model is None:
253+
table_meta = TokenPositionDetails.from_meta(table.this.meta)
254+
table_range_sqlmesh = table_meta.to_range(read_file)
255+
start_pos_sqlmesh = table_range_sqlmesh.start
256+
end_pos_sqlmesh = table_range_sqlmesh.end
257+
references.append(
258+
LSPExternalModelReference(
259+
range=Range(
260+
start=to_lsp_position(start_pos_sqlmesh),
261+
end=to_lsp_position(end_pos_sqlmesh),
262+
),
263+
markdown_description="Unregistered external model",
264+
)
265+
)
250266
continue
251267
referenced_model_path = referenced_model._path
252268
if referenced_model_path is None:

tests/lsp/test_reference.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,9 @@ def test_reference() -> None:
2525
references = get_model_definitions_for_a_path(lsp_context, active_customers_uri)
2626

2727
assert len(references) == 1
28-
assert URI(references[0].uri) == URI.from_path(sushi_customers_path)
28+
uri = references[0].uri
29+
assert uri is not None
30+
assert URI(uri) == URI.from_path(sushi_customers_path)
2931

3032
# Check that the reference in the correct range is sushi.customers
3133
path = active_customers_uri.to_path()

tests/lsp/test_reference_external_model.py

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
1+
from pathlib import Path
2+
13
from lsprotocol.types import Position
4+
5+
from sqlmesh import Config
26
from sqlmesh.core.context import Context
37
from sqlmesh.core.linter.helpers import read_range_from_file
48
from sqlmesh.lsp.context import LSPContext, ModelTarget
59
from sqlmesh.lsp.helpers import to_sqlmesh_range
610
from sqlmesh.lsp.reference import get_references, LSPExternalModelReference
711
from sqlmesh.lsp.uri import URI
12+
from tests.utils.test_filesystem import create_temp_file
813

914

1015
def test_reference() -> None:
@@ -25,14 +30,37 @@ def test_reference() -> None:
2530
assert len(references) == 1
2631
reference = references[0]
2732
assert isinstance(reference, LSPExternalModelReference)
28-
assert reference.uri.endswith("external_models.yaml")
33+
uri = reference.uri
34+
assert uri is not None
35+
assert uri.endswith("external_models.yaml")
2936

3037
source_range = read_range_from_file(customers, to_sqlmesh_range(reference.range))
3138
assert source_range == "raw.demographics"
3239

3340
if reference.target_range is None:
3441
raise AssertionError("Reference target range should not be None")
42+
uri = reference.uri
43+
assert uri is not None
3544
target_range = read_range_from_file(
36-
URI(reference.uri).to_path(), to_sqlmesh_range(reference.target_range)
45+
URI(uri).to_path(), to_sqlmesh_range(reference.target_range)
3746
)
3847
assert target_range == "raw.demographics"
48+
49+
50+
def test_unregistered_external_model(tmp_path: Path):
51+
model_path = tmp_path / "models" / "foo.sql"
52+
contents = "MODEL (name test.foo, kind FULL); SELECT * FROM external_model"
53+
create_temp_file(tmp_path, model_path, contents)
54+
ctx = Context(paths=[tmp_path], config=Config())
55+
lsp_context = LSPContext(ctx)
56+
57+
uri = URI.from_path(model_path)
58+
references = get_references(lsp_context, uri, Position(line=0, character=len(contents) - 3))
59+
60+
assert len(references) == 1
61+
reference = references[0]
62+
assert isinstance(reference, LSPExternalModelReference)
63+
assert reference.uri is None
64+
assert reference.target_range is None
65+
assert reference.markdown_description == "Unregistered external model"
66+
assert read_range_from_file(model_path, to_sqlmesh_range(reference.range)) == "external_model"

0 commit comments

Comments
 (0)