refactor(lsp): move autocomplete functions about (#4603)

Author: benfdking

sqlmesh/lsp/completions.py

@@ -6,7 +6,9 @@
 from sqlmesh.lsp.uri import URI
 
 
-def get_sql_completions(context: t.Optional[LSPContext], file_uri: URI) -> AllModelsResponse:
+def get_sql_completions(
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+) -> AllModelsResponse:
     """
     Return a list of completions for a given file.
     """

sqlmesh/lsp/context.py

@@ -6,6 +6,7 @@
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
 from sqlmesh.lsp.custom import RenderModelEntry, ModelForRendering
+from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
 from sqlmesh.lsp.uri import URI
 
 
@@ -174,3 +175,16 @@ def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
             for audit in self.context.standalone_audits.values()
             if audit._path is not None
         ]
+
+    def get_autocomplete(self, uri: t.Optional[URI]) -> AllModelsResponse:
+        """Get autocomplete suggestions for a file.
+
+        Args:
+            uri: The URI of the file to get autocomplete suggestions for.
+
+        Returns:
+            AllModelsResponse containing models and keywords.
+        """
+        from sqlmesh.lsp.completions import get_sql_completions
+
+        return get_sql_completions(self, uri)

sqlmesh/lsp/main.py

@@ -20,7 +20,6 @@
     ApiResponseGetLineage,
     ApiResponseGetModels,
 )
-from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
     ModelTarget,
@@ -115,9 +114,11 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
             uri = URI(params.textDocument.uri)
             try:
                 context = self._context_get_or_load(uri)
-                return get_sql_completions(context, uri)
+                return context.get_autocomplete(uri)
             except Exception as e:
-                return get_sql_completions(None, uri)
+                from sqlmesh.lsp.completions import get_sql_completions
+
+                return get_sql_completions(None, URI(params.textDocument.uri))
 
         @self.server.feature(RENDER_MODEL_FEATURE)
         def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:

tests/lsp/test_completions.py

@@ -26,7 +26,7 @@ def test_get_sql_completions_with_context_no_file_uri():
     context = Context(paths=["examples/sushi"])
     lsp_context = LSPContext(context)
 
-    completions = get_sql_completions(lsp_context, None)
+    completions = lsp_context.get_autocomplete(None)
     assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
     assert "sushi.active_customers" in completions.models
     assert "sushi.customers" in completions.models
@@ -38,6 +38,6 @@ def test_get_sql_completions_with_context_and_file_uri():
     lsp_context = LSPContext(context)
 
     file_uri = next(key for key in lsp_context.map.keys() if key.name == "active_customers.sql")
-    completions = get_sql_completions(lsp_context, URI.from_path(file_uri))
+    completions = lsp_context.get_autocomplete(URI.from_path(file_uri))
     assert len(completions.keywords) > len(TOKENIZER_KEYWORDS)
     assert "sushi.active_customers" not in completions.models