diff --git a/examples/sushi/models/latest_order.sql b/examples/sushi/models/latest_order.sql index 31293f19f9..4523537505 100644 --- a/examples/sushi/models/latest_order.sql +++ b/examples/sushi/models/latest_order.sql @@ -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 - diff --git a/sqlmesh/lsp/completions.py b/sqlmesh/lsp/completions.py index 7e3781a550..0026260481 100644 --- a/sqlmesh/lsp/completions.py +++ b/sqlmesh/lsp/completions.py @@ -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 @@ -26,14 +31,18 @@ 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. @@ -41,23 +50,23 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t. 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] = [] + 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( diff --git a/sqlmesh/lsp/custom.py b/sqlmesh/lsp/custom.py index 5c8123b7a0..618b4a44bc 100644 --- a/sqlmesh/lsp/custom.py +++ b/sqlmesh/lsp/custom.py @@ -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] diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py index b028ea7766..75ac9b70a2 100755 --- a/sqlmesh/lsp/main.py +++ b/sqlmesh/lsp/main.py @@ -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 diff --git a/tests/lsp/test_completions.py b/tests/lsp/test_completions.py index 7e193d77d6..e0772c1a96 100644 --- a/tests/lsp/test_completions.py +++ b/tests/lsp/test_completions.py @@ -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)