Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion examples/sushi/models/latest_order.sql
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,3 @@ MODEL (
SELECT id, customer_id, start_ts, end_ts, event_date
FROM sushi.orders
ORDER BY event_date DESC LIMIT 1

47 changes: 28 additions & 19 deletions sqlmesh/lsp/completions.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,14 @@
from functools import lru_cache
from sqlglot import Dialect, Tokenizer
from sqlmesh.lsp.custom import AllModelsResponse, MacroCompletion
from sqlmesh.lsp.custom import (
AllModelsResponse,
MacroCompletion,
ModelCompletion,
)
from sqlmesh import macro
import typing as t
from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
from sqlmesh.lsp.description import generate_markdown_description
from sqlmesh.lsp.uri import URI


Expand All @@ -26,38 +31,42 @@ def get_sql_completions(
# Combine keywords - SQL keywords first, then file keywords
all_keywords = list(sql_keywords) + list(file_keywords - sql_keywords)

models = list(get_models(context, file_uri))
return AllModelsResponse(
models=list(get_models(context, file_uri)),
models=[m.name for m in models],
model_completions=models,
keywords=all_keywords,
macros=list(get_macros(context, file_uri)),
)


def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
def get_models(
context: t.Optional[LSPContext], file_uri: t.Optional[URI]
) -> t.List[ModelCompletion]:
"""
Return a list of models for a given file.

If there is no context, return an empty list.
If there is a context, return a list of all models bar the ones the file itself defines.
"""
if context is None:
return set()
return []

current_path = file_uri.to_path() if file_uri is not None else None

completions: t.List[ModelCompletion] = []
Comment thread
benfdking marked this conversation as resolved.
for model in context.context.models.values():
if current_path is not None and model._path == current_path:
continue
description = None
try:
description = generate_markdown_description(model)
except Exception:
description = getattr(model, "description", None)

completions.append(ModelCompletion(name=model.name, description=description))

all_models = set()
# Extract model names from ModelInfo objects
for file_info in context.map.values():
if isinstance(file_info, ModelTarget):
all_models.update(file_info.names)

# Remove models from the current file
path = file_uri.to_path() if file_uri is not None else None
if path is not None and path in context.map:
file_info = context.map[path]
if isinstance(file_info, ModelTarget):
for model in file_info.names:
all_models.discard(model)

return all_models
return completions


def get_macros(
Expand Down
13 changes: 10 additions & 3 deletions sqlmesh/lsp/custom.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,19 @@ class MacroCompletion(PydanticModel):
description: t.Optional[str] = None


class ModelCompletion(PydanticModel):
"""Information about a model for autocompletion."""

name: str
description: t.Optional[str] = None


class AllModelsResponse(CustomMethodResponseBaseClass):
"""
Response to get all the models that are in the current project.
"""
"""Response to get all models that are in the current project."""

#: Deprecated: use ``model_completions`` instead
models: t.List[str]
model_completions: t.List[ModelCompletion]
keywords: t.List[str]
macros: t.List[MacroCompletion]

Expand Down
10 changes: 8 additions & 2 deletions sqlmesh/lsp/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -627,12 +627,18 @@ def completion(

completion_items = []
# Add model completions
for model in completion_response.models:
for model in completion_response.model_completions:
completion_items.append(
types.CompletionItem(
label=model,
label=model.name,
kind=types.CompletionItemKind.Reference,
detail="SQLMesh Model",
documentation=types.MarkupContent(
kind=types.MarkupKind.Markdown,
value=model.description or "No description available",
)
if model.description
else None,
)
)
# Add macro completions
Expand Down
14 changes: 14 additions & 0 deletions tests/lsp/test_completions.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,20 @@ def test_get_macros():
assert add_one_macro.description


def test_model_completions_include_descriptions():
context = Context(paths=["examples/sushi"])
lsp_context = LSPContext(context)

completions = LSPContext.get_completions(lsp_context, None)

model_entry = next(
(m for m in completions.model_completions if m.name == "sushi.customers"),
None,
)
assert model_entry is not None
assert model_entry.description


def test_get_sql_completions_with_context_no_file_uri():
context = Context(paths=["examples/sushi"])
lsp_context = LSPContext(context)
Expand Down