faet(vscode): support multiple workspaces (#4634)

Author: benfdking

sqlmesh/lsp/main.py

@@ -56,6 +56,7 @@ def __init__(
         self.server = LanguageServer(server_name, version)
         self.context_class = context_class
         self.lsp_context: t.Optional[LSPContext] = None
+        self.workspace_folders: t.List[Path] = []
 
         self.client_supports_pull_diagnostics = False
         # Register LSP features (e.g., formatting, hover, etc.)
@@ -68,7 +69,7 @@ def _register_features(self) -> None:
         def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
             """Initialize the server when the client connects."""
             try:
-                # Check if client supports pull diagnostics
+                # Check if the client supports pull diagnostics
                 if params.capabilities and params.capabilities.text_document:
                     diagnostics = getattr(params.capabilities.text_document, "diagnostic", None)
                     if diagnostics:
@@ -81,9 +82,13 @@ def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
                     self.client_supports_pull_diagnostics = False
 
                 if params.workspace_folders:
+                    # Store all workspace folders for later use
+                    self.workspace_folders = [
+                        Path(self._uri_to_path(folder.uri)) for folder in params.workspace_folders
+                    ]
+
                     # Try to find a SQLMesh config file in any workspace folder (only at the root level)
-                    for folder in params.workspace_folders:
-                        folder_path = Path(self._uri_to_path(folder.uri))
+                    for folder_path in self.workspace_folders:
                         # Only check for config files directly in the workspace directory
                         for ext in ("py", "yml", "yaml"):
                             config_path = folder_path / f"config.{ext}"
@@ -104,8 +109,8 @@ 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)
             try:
-                uri = URI(params.textDocument.uri)
                 context = self._context_get_or_load(uri)
                 return get_sql_completions(context, uri)
             except Exception as e:
@@ -457,6 +462,8 @@ def _context_get_or_load(self, document_uri: URI) -> LSPContext:
     def _ensure_context_in_folder(self, folder_uri: Path) -> None:
         if self.lsp_context is not None:
             return
+
+        # First, check the provided folder
         for ext in ("py", "yml", "yaml"):
             config_path = folder_uri / f"config.{ext}"
             if config_path.exists():
@@ -468,6 +475,22 @@ def _ensure_context_in_folder(self, folder_uri: Path) -> None:
                 except Exception as e:
                     self.server.show_message(f"Error loading context: {e}", types.MessageType.Error)
 
+        # If not found in the provided folder, search through all workspace folders
+        for workspace_folder in self.workspace_folders:
+            for ext in ("py", "yml", "yaml"):
+                config_path = workspace_folder / f"config.{ext}"
+                if config_path.exists():
+                    try:
+                        created_context = self.context_class(paths=[workspace_folder])
+                        self.lsp_context = LSPContext(created_context)
+                        loaded_sqlmesh_message(self.server, workspace_folder)
+                        return
+                    except Exception as e:
+                        self.server.show_message(
+                            f"Error loading context from {config_path}: {e}",
+                            types.MessageType.Warning,
+                        )
+
     def _ensure_context_for_document(
         self,
         document_uri: URI,
@@ -506,6 +529,23 @@ def _ensure_context_for_document(
                         )
             path = path.parent
 
+        # If still no context found, try the workspace folders
+        if not loaded:
+            for workspace_folder in self.workspace_folders:
+                for ext in ("py", "yml", "yaml"):
+                    config_path = workspace_folder / f"config.{ext}"
+                    if config_path.exists():
+                        try:
+                            created_context = self.context_class(paths=[workspace_folder])
+                            self.lsp_context = LSPContext(created_context)
+                            loaded_sqlmesh_message(self.server, workspace_folder)
+                            return
+                        except Exception as e:
+                            self.server.show_message(
+                                f"Error loading context from {config_path}: {e}",
+                                types.MessageType.Warning,
+                            )
+
         return
 
     @staticmethod

vscode/extension/src/lsp/lsp.ts

@@ -34,13 +34,11 @@ export class LSPClient implements Disposable {
       return sqlmesh
     }
     const workspaceFolders = getWorkspaceFolders()
-    if (workspaceFolders.length !== 1) {
-      traceError(
-        `Invalid number of workspace folders: ${workspaceFolders.length}`,
-      )
+    if (workspaceFolders.length === 0) {
+      traceError(`No workspace folders found`)
       return err({
         type: 'generic',
-        message: 'Invalid number of workspace folders',
+        message: 'No workspace folders found',
       })
     }
     const workspacePath = sqlmesh.value.workspacePath

vscode/extension/tests/lineage.spec.ts

@@ -3,6 +3,7 @@ import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
 import { startVSCode, SUSHI_SOURCE_PATH } from './utils'
+import { writeFileSync } from 'fs'
 
 /**
  * Helper function to launch VS Code and test lineage with given project path config
@@ -15,8 +16,8 @@ async function testLineageWithProjectPath(window: Page): Promise<void> {
   await window.keyboard.type('Lineage: Focus On View')
   await window.keyboard.press('Enter')
 
-  // Wait for "Loaded SQLmesh Context" text to appear
-  const loadedContextText = window.locator('text=Loaded SQLMesh Context')
+  // Wait for "Loaded SQLMesh context" text to appear
+  const loadedContextText = window.locator('text=Loaded SQLMesh context')
   await expect(loadedContextText.first()).toBeVisible({ timeout: 10_000 })
 }
 
@@ -141,3 +142,81 @@ test('Lineage panel renders correctly - absolute path project outside of workspa
     await fs.remove(tempFolder)
   }
 })
+
+test('Lineage panel renders correctly - multiworkspace setup', async () => {
+  const workspaceDir = await fs.mkdtemp(
+    path.join(os.tmpdir(), 'vscode-test-workspace-'),
+  )
+  const projectDir1 = path.join(workspaceDir, 'projects', 'sushi1')
+  const projectDir2 = path.join(workspaceDir, 'projects', 'sushi2')
+  await fs.copy(SUSHI_SOURCE_PATH, projectDir1)
+  await fs.ensureDir(projectDir2)
+
+  // Add a .code-workspace file with multiple projects
+  const workspaceFilePath = path.join(
+    workspaceDir,
+    'multi-workspace.code-workspace',
+  )
+  writeFileSync(
+    workspaceFilePath,
+    JSON.stringify({
+      folders: [
+        {
+          name: 'sushi1',
+          path: 'projects/sushi1',
+        },
+        {
+          name: 'sushi2',
+          path: 'projects/sushi2',
+        },
+      ],
+    }),
+  )
+
+  try {
+    const { window, close } = await startVSCode(workspaceFilePath)
+    await testLineageWithProjectPath(window)
+    await close()
+  } finally {
+    await fs.remove(workspaceDir)
+  }
+})
+
+test('Lineage panel renders correctly - multiworkspace setup reversed', async () => {
+  const workspaceDir = await fs.mkdtemp(
+    path.join(os.tmpdir(), 'vscode-test-workspace-'),
+  )
+  const projectDir1 = path.join(workspaceDir, 'projects', 'sushi1')
+  const projectDir2 = path.join(workspaceDir, 'projects', 'sushi2')
+  await fs.copy(SUSHI_SOURCE_PATH, projectDir2)
+  await fs.ensureDir(projectDir1)
+
+  // Add a .code-workspace file with multiple projects
+  const workspaceFilePath = path.join(
+    workspaceDir,
+    'multi-workspace.code-workspace',
+  )
+  writeFileSync(
+    workspaceFilePath,
+    JSON.stringify({
+      folders: [
+        {
+          name: 'sushi1',
+          path: 'projects/sushi1',
+        },
+        {
+          name: 'sushi2',
+          path: 'projects/sushi2',
+        },
+      ],
+    }),
+  )
+
+  try {
+    const { window, close } = await startVSCode(workspaceFilePath)
+    await testLineageWithProjectPath(window)
+    await close()
+  } finally {
+    await fs.remove(workspaceDir)
+  }
+})