feat(lsp): add keywords in query to autocompete

benfdking · benfdking · commit 6b9cbf3bde2b · 2025-06-05T09:52:56.000+01:00
diff --git a/sqlmesh/lsp/completions.py b/sqlmesh/lsp/completions.py
@@ -7,14 +7,25 @@
 
 
 def get_sql_completions(
-    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI], content: t.Optional[str] = None
 ) -> AllModelsResponse:
     """
     Return a list of completions for a given file.
     """
+    # Get SQL keywords for the dialect
+    sql_keywords = get_keywords(context, file_uri)
+
+    # Get keywords from file content if provided
+    file_keywords = set()
+    if content:
+        file_keywords = extract_keywords_from_content(content, get_dialect(context, file_uri))
+
+    # Combine keywords - SQL keywords first, then file keywords
+    all_keywords = list(sql_keywords) + list(file_keywords - sql_keywords)
+
     return AllModelsResponse(
         models=list(get_models(context, file_uri)),
-        keywords=list(get_keywords(context, file_uri)),
+        keywords=all_keywords,
     )
 
 
@@ -97,3 +108,54 @@ def get_keywords_from_tokenizer(dialect: t.Optional[str] = None) -> t.Set[str]:
         parts = keyword.split(" ")
         expanded_keywords.update(parts)
     return expanded_keywords
+
+
+def get_dialect(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Optional[str]:
+    """
+    Get the dialect for a given file.
+    """
+    if file_uri is not None and context is not None and file_uri.to_path() in context.map:
+        file_info = context.map[file_uri.to_path()]
+
+        # Handle ModelInfo objects
+        if isinstance(file_info, ModelTarget) and file_info.names:
+            model_name = file_info.names[0]
+            model_from_context = context.context.get_model(model_name)
+            return model_from_context.dialect
+
+        # Handle AuditInfo objects
+        if isinstance(file_info, AuditTarget) and file_info.name:
+            audit = context.context.standalone_audits.get(file_info.name)
+            if audit is not None and audit.dialect:
+                return audit.dialect
+
+    if context is not None:
+        return context.context.default_dialect
+
+    return None
+
+
+def extract_keywords_from_content(content: str, dialect: t.Optional[str] = None) -> t.Set[str]:
+    """
+    Extract identifiers from SQL content using the tokenizer.
+    Only extracts identifiers (variable names, table names, column names, etc.)
+    that are not SQL keywords.
+    """
+    if not content:
+        return set()
+
+    tokenizer_class = Dialect.get_or_raise(dialect).tokenizer_class
+    keywords = set()
+    try:
+        tokenizer = tokenizer_class()
+        tokens = tokenizer.tokenize(content)
+        for token in tokens:
+            # Don't include keywords in the set
+            if token.text.upper() not in tokenizer_class.KEYWORDS:
+                keywords.add(token.text)
+
+    except Exception:
+        # If tokenization fails, return empty set
+        pass
+
+    return keywords
diff --git a/sqlmesh/lsp/context.py b/sqlmesh/lsp/context.py
@@ -176,15 +176,18 @@ def list_of_models_for_rendering(self) -> t.List[ModelForRendering]:
             if audit._path is not None
         ]
 
-    def get_autocomplete(self, uri: t.Optional[URI]) -> AllModelsResponse:
+    def get_autocomplete(
+        self, uri: t.Optional[URI], content: t.Optional[str] = None
+    ) -> AllModelsResponse:
         """Get autocomplete suggestions for a file.
 
         Args:
             uri: The URI of the file to get autocomplete suggestions for.
+            content: The content of the file (optional).
 
         Returns:
             AllModelsResponse containing models and keywords.
         """
         from sqlmesh.lsp.completions import get_sql_completions
 
-        return get_sql_completions(self, uri)
+        return get_sql_completions(self, uri, content)
diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py
@@ -114,13 +114,22 @@ def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
         @self.server.feature(ALL_MODELS_FEATURE)
         def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsResponse:
             uri = URI(params.textDocument.uri)
+
+            # Get the document content
+            content = None
+            try:
+                document = ls.workspace.get_text_document(params.textDocument.uri)
+                content = document.source
+            except Exception:
+                pass
+
             try:
                 context = self._context_get_or_load(uri)
-                return context.get_autocomplete(uri)
+                return context.get_autocomplete(uri, content)
             except Exception as e:
                 from sqlmesh.lsp.completions import get_sql_completions
 
-                return get_sql_completions(None, URI(params.textDocument.uri))
+                return get_sql_completions(None, URI(params.textDocument.uri), content)
 
         @self.server.feature(RENDER_MODEL_FEATURE)
         def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:
@@ -471,8 +480,16 @@ def completion(
                 uri = URI(params.text_document.uri)
                 context = self._context_get_or_load(uri)
 
+                # Get the document content
+                content = None
+                try:
+                    document = ls.workspace.get_text_document(params.text_document.uri)
+                    content = document.source
+                except Exception:
+                    pass
+
                 # Get completions using the existing completions module
-                completion_response = context.get_autocomplete(uri)
+                completion_response = context.get_autocomplete(uri, content)
 
                 completion_items = []
                 # Add model completions
diff --git a/tests/lsp/test_completions.py b/tests/lsp/test_completions.py
@@ -1,7 +1,11 @@
 import pytest
 from sqlglot import Tokenizer
 from sqlmesh.core.context import Context
-from sqlmesh.lsp.completions import get_keywords_from_tokenizer, get_sql_completions
+from sqlmesh.lsp.completions import (
+    get_keywords_from_tokenizer,
+    get_sql_completions,
+    extract_keywords_from_content,
+)
 from sqlmesh.lsp.context import LSPContext
 from sqlmesh.lsp.uri import URI
 
@@ -41,3 +45,131 @@ def test_get_sql_completions_with_context_and_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
+
+
+@pytest.mark.fast
+def test_extract_keywords_from_content():
+    # Test extracting keywords from SQL content
+    content = """
+    SELECT customer_id, order_date, total_amount
+    FROM orders o
+    JOIN customers c ON o.customer_id = c.id
+    WHERE order_date > '2024-01-01'
+    """
+
+    keywords = extract_keywords_from_content(content)
+
+    # Check that identifiers are extracted
+    assert "customer_id" in keywords
+    assert "order_date" in keywords
+    assert "total_amount" in keywords
+    assert "orders" in keywords
+    assert "customers" in keywords
+    assert "o" in keywords  # alias
+    assert "c" in keywords  # alias
+    assert "id" in keywords
+
+    # Check that SQL keywords are NOT included
+    assert "SELECT" not in keywords
+    assert "FROM" not in keywords
+    assert "JOIN" not in keywords
+    assert "WHERE" not in keywords
+    assert "ON" not in keywords
+
+
+@pytest.mark.fast
+def test_get_sql_completions_with_file_content():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # SQL content with custom identifiers
+    content = """
+    SELECT my_custom_column, another_identifier
+    FROM my_custom_table mct
+    JOIN some_other_table sot ON mct.id = sot.table_id
+    WHERE my_custom_column > 100
+    """
+
+    file_uri = next(key for key in lsp_context.map.keys() if key.name == "active_customers.sql")
+    completions = lsp_context.get_autocomplete(URI.from_path(file_uri), content)
+
+    # Check that SQL keywords are included
+    assert any(k in ["SELECT", "FROM", "WHERE", "JOIN"] for k in completions.keywords)
+
+    # Check that file-specific identifiers are included at the end
+    keywords_list = completions.keywords
+    assert "my_custom_column" in keywords_list
+    assert "another_identifier" in keywords_list
+    assert "my_custom_table" in keywords_list
+    assert "some_other_table" in keywords_list
+    assert "mct" in keywords_list  # alias
+    assert "sot" in keywords_list  # alias
+    assert "table_id" in keywords_list
+
+    # Check that file keywords come after SQL keywords
+    # SQL keywords should appear first in the list
+    sql_keyword_indices = [
+        i for i, k in enumerate(keywords_list) if k in ["SELECT", "FROM", "WHERE", "JOIN"]
+    ]
+    file_keyword_indices = [
+        i for i, k in enumerate(keywords_list) if k in ["my_custom_column", "my_custom_table"]
+    ]
+
+    if sql_keyword_indices and file_keyword_indices:
+        assert max(sql_keyword_indices) < min(file_keyword_indices), (
+            "SQL keywords should come before file keywords"
+        )
+
+
+@pytest.mark.fast
+def test_get_sql_completions_with_partial_cte_query():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Partial SQL query with CTEs
+    content = """
+    WITH _latest_complete_month AS (
+        SELECT MAX(date_trunc('month', order_date)) as month
+        FROM orders
+    ),
+    _filtered AS (
+        SELECT * FROM
+    """
+
+    file_uri = next(key for key in lsp_context.map.keys() if key.name == "active_customers.sql")
+    completions = lsp_context.get_autocomplete(URI.from_path(file_uri), content)
+
+    # Check that CTE names are included in the keywords
+    keywords_list = completions.keywords
+    assert "_latest_complete_month" in keywords_list
+    assert "_filtered" in keywords_list
+
+    # Also check other identifiers from the partial query
+    assert "month" in keywords_list
+    assert "order_date" in keywords_list
+    assert "orders" in keywords_list
+
+
+@pytest.mark.fast
+def test_extract_keywords_from_partial_query():
+    # Test extracting keywords from an incomplete SQL query
+    content = """
+    WITH cte1 AS (
+        SELECT col1, col2 FROM table1
+    ),
+    cte2 AS (
+        SELECT * FROM cte1 WHERE
+    """
+
+    keywords = extract_keywords_from_content(content)
+
+    # Check that CTEs are extracted
+    assert "cte1" in keywords
+    assert "cte2" in keywords
+
+    # Check that columns and tables are extracted
+    assert "col1" in keywords
+    assert "col2" in keywords
+    assert "table1" in keywords
+
+    # Even though the query is incomplete, identifiers should still be extracted
