fix(vscode): ensure only load context once

Author: benfdking

sqlmesh/lsp/main.py

@@ -85,7 +85,7 @@ class ContextFailed:
     """State when context failed to load with an error message."""
 
     error: ContextFailedError
-    context: t.Optional[Context] = None
+    context: Context
     version_id: str = field(default_factory=lambda: str(uuid.uuid4()))
 
 
@@ -227,24 +227,26 @@ def _reload_context_and_publish_diagnostics(
         self, ls: LanguageServer, uri: URI, document_uri: str
     ) -> None:
         """Helper method to reload context and publish diagnostics."""
+        context: t.Optional[Context] = None
         if isinstance(self.context_state, NoContext):
             pass
         elif isinstance(self.context_state, ContextFailed):
             if self.context_state.context:
+                context = self.context_state.context
                 try:
-                    self.context_state.context.load()
+                    context.load()
                     # Creating a new LSPContext will naturally create fresh caches
-                    self.context_state = ContextLoaded(
-                        lsp_context=LSPContext(self.context_state.context)
-                    )
+                    self.context_state = ContextLoaded(lsp_context=LSPContext(context))
                 except Exception as e:
                     ls.log_trace(f"Error loading context: {e}")
-                    context = (
+                    error_context = (
                         self.context_state.context
                         if hasattr(self.context_state, "context")
                         else None
                     )
-                    self.context_state = ContextFailed(error=e, context=context)
+                    if error_context is None:
+                        raise RuntimeError(f"Context should always be set.")
+                    self.context_state = ContextFailed(error=e, context=error_context)
             else:
                 # If there's no context, reset to NoContext and try to create one from scratch
                 ls.log_trace("No partial context available, attempting fresh creation")
@@ -860,17 +862,20 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
         Returns:
             A new LSPContext instance wrapping the created context, or None if creation fails
         """
+        context = None
         try:
             if isinstance(self.context_state, NoContext):
-                context = self.context_class(paths=paths)
+                context = self.context_class(paths=paths, load=False)
+                context.load()
                 loaded_sqlmesh_message(self.server, paths[0])
             elif isinstance(self.context_state, ContextFailed):
                 if self.context_state.context:
                     context = self.context_state.context
                     context.load()
                 else:
                     # If there's no context (initial creation failed), try creating again
-                    context = self.context_class(paths=paths)
+                    context = self.context_class(paths=paths, load=False)
+                    context.load()
                     loaded_sqlmesh_message(self.server, paths[0])
             else:
                 context = self.context_state.lsp_context.context
@@ -889,11 +894,12 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
             self.server.log_trace(f"Error creating context: {e}")
             # Store the error in context state so subsequent requests show the actual error
             # Try to preserve any partially loaded context if it exists
-            context = None
             if isinstance(self.context_state, ContextLoaded):
                 context = self.context_state.lsp_context.context
             elif isinstance(self.context_state, ContextFailed) and self.context_state.context:
                 context = self.context_state.context
+            if context is None:
+                raise RuntimeError("Context should always be set.")
             self.context_state = ContextFailed(error=e, context=context)
             return None
 

vscode/extension/tests/broken_project.spec.ts

@@ -2,7 +2,12 @@ import { test, expect } from './fixtures'
 import fs from 'fs-extra'
 import os from 'os'
 import path from 'path'
-import { openLineageView, saveFile, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openLineageView,
+  runCommand,
+  saveFile,
+  SUSHI_SOURCE_PATH,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('bad project, double model', async ({ page, sharedCodeServer }) => {
@@ -253,3 +258,54 @@ test('bad project, double model, check lineage', async ({
 
   await page.waitForTimeout(500)
 })
+
+test('bad model block, then fixed', async ({ page, sharedCodeServer }) => {
+  const tempDir = await fs.mkdtemp(
+    path.join(os.tmpdir(), 'vscode-test-tcloud-'),
+  )
+  // Copy over the sushi project
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+  await createPythonInterpreterSettingsSpecifier(tempDir)
+
+  // Add a model with a bad model block
+  const badModelPath = path.join(tempDir, 'models', 'bad_model.sql')
+  const contents =
+    'MODEL ( name sushi.bad_block, test); SELECT * FROM sushi.customers'
+  await fs.writeFile(badModelPath, contents)
+
+  await page.goto(
+    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
+  )
+  await page.waitForLoadState('networkidle')
+
+  // Open the customers.sql model
+  await page
+    .getByRole('treeitem', { name: 'models', exact: true })
+    .locator('a')
+    .click()
+  await page
+    .getByRole('treeitem', { name: 'customers.sql', exact: true })
+    .locator('a')
+    .click()
+
+  // Wait for the error to appear
+  await page.waitForSelector('text=Error creating context')
+
+  // Open the problems view
+  await runCommand(page, 'View: Focus Problems')
+
+  // Assert error is present in the problems view
+  const errorElement = page
+    .getByText("Required keyword: 'value' missing for")
+    .first()
+  await expect(errorElement).toBeVisible({ timeout: 5000 })
+
+  // Remove the bad model file
+  await fs.remove(badModelPath)
+
+  // Save to refresh the context
+  await saveFile(page)
+
+  // Wait for successful context load
+  await page.waitForSelector('text=Loaded SQLMesh context')
+})