chore(lsp): move autocomplete to the lsp

Author: benfdking

sqlmesh/lsp/main.py

@@ -20,6 +20,7 @@
     ApiResponseGetLineage,
     ApiResponseGetModels,
 )
+from sqlmesh.lsp.completions import get_sql_completions
 from sqlmesh.lsp.context import (
     LSPContext,
     ModelTarget,
@@ -456,6 +457,47 @@ def workspace_diagnostic(
                 )
                 return types.WorkspaceDiagnosticReport(items=[])
 
+        @self.server.feature(types.TEXT_DOCUMENT_COMPLETION)
+        def completion(
+            ls: LanguageServer, params: types.CompletionParams
+        ) -> t.Optional[types.CompletionList]:
+            """Handle completion requests from the client."""
+            try:
+                uri = URI(params.text_document.uri)
+                context = self._context_get_or_load(uri)
+
+                # Get completions using the existing completions module
+                completion_response = context.get_autocomplete(uri)
+
+                completion_items = []
+                # Add model completions
+                for model in completion_response.models:
+                    completion_items.append(
+                        types.CompletionItem(
+                            label=model,
+                            kind=types.CompletionItemKind.Reference,
+                            detail="SQLMesh Model",
+                        )
+                    )
+                # Add keyword completions
+                for keyword in completion_response.keywords:
+                    completion_items.append(
+                        types.CompletionItem(
+                            label=keyword,
+                            kind=types.CompletionItemKind.Keyword,
+                            detail="SQL Keyword",
+                        )
+                    )
+
+                return types.CompletionList(
+                    is_incomplete=False,
+                    items=completion_items,
+                )
+
+            except Exception as e:
+                get_sql_completions(None, URI(params.text_document.uri))
+                return None
+
     def _get_diagnostics_for_uri(self, uri: URI) -> t.Tuple[t.List[types.Diagnostic], int]:
         """Get diagnostics for a specific URI, returning (diagnostics, result_id).
 

vscode/extension/src/extension.ts

@@ -84,13 +84,6 @@ export async function activate(context: vscode.ExtensionContext) {
     ),
   )
 
-  context.subscriptions.push(
-    vscode.languages.registerCompletionItemProvider(
-      selector,
-      completionProvider(lspClient),
-    ),
-  )
-
   // Register the webview
   const lineagePanel = new LineagePanel(context.extensionUri, lspClient)
   context.subscriptions.push(
@@ -185,6 +178,15 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(lspClient)
   }
 
+  if (lspClient && !lspClient.hasCompletionCapability()) {
+    context.subscriptions.push(
+      vscode.languages.registerCompletionItemProvider(
+        selector,
+        completionProvider(lspClient),
+      ),
+    )
+  }
+
   traceInfo('Extension activated')
 }
 

vscode/extension/src/lsp/lsp.ts

@@ -21,6 +21,16 @@ export class LSPClient implements Disposable {
     this.client = undefined
   }
 
+  public hasCompletionCapability(): boolean {
+    if (!this.client) {
+      traceError('LSP client is not initialized')
+      return false
+    }
+    const capabilities = this.client.initializeResult?.capabilities
+    const completion = capabilities?.completionProvider
+    return completion !== undefined
+  }
+
   public async start(): Promise<Result<undefined, ErrorType>> {
     if (!outputChannel) {
       outputChannel = window.createOutputChannel('sqlmesh-lsp')

vscode/extension/tests/completions.spec.ts

@@ -38,13 +38,15 @@ test('Autocomplete for model names', async () => {
     await window.keyboard.press('Enter')
 
     // Type the beginning of sushi.customers to trigger autocomplete
-    await window.keyboard.type('sushi.cus')
+    await window.keyboard.type('sushi.waiter_as_customer')
 
     // Wait a moment for autocomplete to appear
     await window.waitForTimeout(500)
 
     // Check if the autocomplete suggestion for sushi.customers is visible
-    await expect(window.locator('text=sushi.customers')).toBeVisible()
+    expect(
+      await window.locator('text=sushi.waiter_as_customer_by_day').count(),
+    ).toBe(1)
 
     await close()
   } finally {