feat(vscode): multi project for vscode (#5084)

Author: benfdking

.gitignore

@@ -157,4 +157,4 @@ metastore_db/
 spark-warehouse/
 
 # claude
-.claude/
+.claude/

examples/multi/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "sqlmesh.projectPaths": ["./repo_1", "./repo_2"]
+}

sqlmesh/lsp/main.py

@@ -71,12 +71,20 @@
 from sqlmesh.lsp.uri import URI
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.lineage import ExternalModelReference
+from sqlmesh.utils.pydantic import PydanticModel
 from web.server.api.endpoints.lineage import column_lineage, model_lineage
 from web.server.api.endpoints.models import get_models
 from typing import Union
 from dataclasses import dataclass, field
 
 
+class InitializationOptions(PydanticModel):
+    """Initialization options for the SQLMesh Language Server, that
+    are passed from the client to the server."""
+
+    project_paths: t.Optional[t.List[str]] = None
+
+
 @dataclass
 class NoContext:
     """State when no context has been attempted to load."""
@@ -105,6 +113,11 @@ class ContextFailed:
 
 
 class SQLMeshLanguageServer:
+    # Specified folders take precedence over workspace folders or looking
+    # for a config files. They are explicitly set by the user and optionally
+    # pass in at init
+    specified_paths: t.Optional[t.List[Path]] = None
+
     def __init__(
         self,
         context_class: t.Type[Context],
@@ -411,6 +424,12 @@ def command_external_models_update_columns(ls: LanguageServer, raw: t.Any) -> No
         def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
             """Initialize the server when the client connects."""
             try:
+                # Check the custom options
+                if params.initialization_options:
+                    options = InitializationOptions.model_validate(params.initialization_options)
+                    if options.project_paths is not None:
+                        self.specified_paths = [Path(path) for path in options.project_paths]
+
                 # Check if the client supports pull diagnostics
                 if params.capabilities and params.capabilities.text_document:
                     diagnostics = getattr(params.capabilities.text_document, "diagnostic", None)
@@ -906,7 +925,12 @@ def _context_get_or_load(self, document_uri: t.Optional[URI] = None) -> LSPConte
                 raise Exception(state.error)
             raise state.error
         if isinstance(state, NoContext):
-            self._ensure_context_for_document(document_uri)
+            if self.specified_paths is not None:
+                # If specified paths are provided, create context from them
+                if self._create_lsp_context(self.specified_paths):
+                    loaded_sqlmesh_message(self.server)
+            else:
+                self._ensure_context_for_document(document_uri)
         if isinstance(state, ContextLoaded):
             return state.lsp_context
         raise RuntimeError("Context failed to load")

vscode/extension/package.json

@@ -32,10 +32,13 @@
       "type": "object",
       "title": "SQLMesh",
       "properties": {
-        "sqlmesh.projectPath": {
-          "type": "string",
-          "default": "",
-          "markdownDescription": "The path to the SQLMesh project. If not set, the extension will try to find the project root automatically. If set, the extension will use the project root as the workspace path, e.g. it will run `sqlmesh` and `sqlmesh_lsp` in the project root. The path can be absolute `/Users/sqlmesh_user/sqlmesh_project/sushi` or relative `./project_folder/sushi` to the workspace root."
+        "sqlmesh.projectPaths": {
+          "type": "array",
+          "items": {
+            "type": "string"
+          },
+          "default": [],
+          "description": "The path to the SQLMesh project. If not set, the extension will try to find the project root automatically. If set, the extension will use the project root as the workspace path, e.g. it will run `sqlmesh` and `sqlmesh_lsp` in the project root. The path can be absolute `/Users/sqlmesh_user/sqlmesh_project/sushi` or relative `./project_folder/sushi` to the workspace root. Multiple paths can be used for multi-project setups."
         },
         "sqlmesh.lspEntrypoint": {
           "type": "string",

vscode/extension/src/lsp/lsp.ts

@@ -16,6 +16,7 @@ import {
   ErrorTypeSQLMeshOutdated,
 } from '../utilities/errors'
 import { CustomLSPMethods } from './custom'
+import { resolveProjectPath } from '../utilities/config'
 
 type SupportedMethodsState =
   | { type: 'not-fetched' }
@@ -109,6 +110,11 @@ export class LSPClient implements Disposable {
         args: sqlmesh.value.args,
       },
     }
+    const paths = resolveProjectPath(getWorkspaceFolders()[0])
+    if (isErr(paths)) {
+      traceError(`Failed to resolve project paths: ${paths.error}`)
+      return err({ type: 'generic', message: paths.error })
+    }
     const clientOptions: LanguageClientOptions = {
       documentSelector: [
         { scheme: 'file', pattern: '**/*.sql' },
@@ -117,6 +123,11 @@ export class LSPClient implements Disposable {
       ],
       diagnosticCollectionName: 'sqlmesh',
       outputChannel,
+      initializationOptions: paths.value.projectPaths
+        ? {
+            project_paths: paths.value.projectPaths,
+          }
+        : null,
     }
 
     traceInfo(

vscode/extension/src/utilities/config.ts

@@ -1,17 +1,17 @@
 import { workspace, WorkspaceFolder } from 'vscode'
 import path from 'path'
 import fs from 'fs'
-import { Result, err, ok } from '@bus/result'
+import { Result, err, isErr, ok } from '@bus/result'
 import { traceVerbose, traceInfo } from './common/log'
 import { parse } from 'shell-quote'
 import { z } from 'zod'
 
-const configSchema = z.object({
-  projectPath: z.string(),
+const sqlmeshConfigurationSchema = z.object({
+  projectPaths: z.array(z.string()),
   lspEntryPoint: z.string(),
 })
 
-export type SqlmeshConfiguration = z.infer<typeof configSchema>
+export type SqlmeshConfiguration = z.infer<typeof sqlmeshConfigurationSchema>
 
 /**
  * Get the SQLMesh configuration from VS Code settings.
@@ -20,16 +20,15 @@ export type SqlmeshConfiguration = z.infer<typeof configSchema>
  */
 function getSqlmeshConfiguration(): SqlmeshConfiguration {
   const config = workspace.getConfiguration('sqlmesh')
-  const projectPath = config.get<string>('projectPath', '')
+  const projectPaths = config.get<string[]>('projectPaths', [])
   const lspEntryPoint = config.get<string>('lspEntrypoint', '')
-
-  const parsed = configSchema.safeParse({
-    projectPath,
+  const parsed = sqlmeshConfigurationSchema.safeParse({
+    projectPaths,
     lspEntryPoint,
   })
   if (!parsed.success) {
     throw new Error(
-      `Invalid sqlmesh configuration: ${JSON.stringify(parsed.error)}`,
+      `Invalid SQLMesh configuration: ${JSON.stringify(parsed.error)}`,
     )
   }
   return parsed.data
@@ -66,31 +65,57 @@ export function getSqlmeshLspEntryPoint():
 }
 
 /**
- * Validate and resolve the project path from configuration.
+ * Validate and resolve the project paths from configuration.
  * If no project path is configured, use the workspace folder.
  * If the project path is configured, it must be a directory that contains a SQLMesh project.
  *
  * @param workspaceFolder The current workspace folder
- * @returns A Result containing the resolved project path or an error
+ * @returns A Result containing the resolved project paths or an error
  */
-export function resolveProjectPath(
-  workspaceFolder: WorkspaceFolder,
-): Result<string, string> {
+export function resolveProjectPath(workspaceFolder: WorkspaceFolder): Result<
+  {
+    projectPaths: string[] | undefined
+    workspaceFolder: string
+  },
+  string
+> {
   const config = getSqlmeshConfiguration()
 
-  if (!config.projectPath) {
+  if (config.projectPaths.length === 0) {
     // If no project path is configured, use the workspace folder
     traceVerbose('No project path configured, using workspace folder')
-    return ok(workspaceFolder.uri.fsPath)
+    return ok({
+      workspaceFolder: workspaceFolder.uri.fsPath,
+      projectPaths: undefined,
+    })
+  }
+
+  const resolvedPaths: string[] = []
+  for (const projectPath of config.projectPaths) {
+    const result = resolveSingleProjectPath(workspaceFolder, projectPath)
+    if (isErr(result)) {
+      return result
+    }
+    resolvedPaths.push(result.value)
   }
+  return ok({
+    projectPaths: resolvedPaths,
+    workspaceFolder: workspaceFolder.uri.fsPath,
+  })
+}
+
+function resolveSingleProjectPath(
+  workspaceFolder: WorkspaceFolder,
+  projectPath: string,
+): Result<string, string> {
   let resolvedPath: string
 
   // Check if the path is absolute
-  if (path.isAbsolute(config.projectPath)) {
-    resolvedPath = config.projectPath
+  if (path.isAbsolute(projectPath)) {
+    resolvedPath = projectPath
   } else {
     // Resolve relative path from workspace root
-    resolvedPath = path.join(workspaceFolder.uri.fsPath, config.projectPath)
+    resolvedPath = path.join(workspaceFolder.uri.fsPath, projectPath)
   }
 
   // Normalize the path

vscode/extension/src/utilities/sqlmesh/sqlmesh.ts

@@ -69,8 +69,8 @@ export const isTcloudProject = async (): Promise<Result<boolean, string>> => {
   if (isErr(resolvedPath)) {
     return err(resolvedPath.error)
   }
-  const tcloudYamlPath = path.join(resolvedPath.value, 'tcloud.yaml')
-  const tcloudYmlPath = path.join(resolvedPath.value, 'tcloud.yml')
+  const tcloudYamlPath = path.join(resolvedPath.value.workspaceFolder, 'tcloud.yaml')
+  const tcloudYmlPath = path.join(resolvedPath.value.workspaceFolder, 'tcloud.yml')
   const isTcloudYamlFilePresent = fs.existsSync(tcloudYamlPath)
   const isTcloudYmlFilePresent = fs.existsSync(tcloudYmlPath)
   if (isTcloudYamlFilePresent || isTcloudYmlFilePresent) {
@@ -144,7 +144,7 @@ export const isSqlmeshEnterpriseInstalled = async (): Promise<
     })
   }
   const called = await execAsync(tcloudBin.value.bin, ['is_sqlmesh_installed'], {
-    cwd: resolvedPath.value,
+    cwd: resolvedPath.value.workspaceFolder,
     env: tcloudBin.value.env,
   })
   if (called.exitCode !== 0) {
@@ -185,7 +185,7 @@ export const installSqlmeshEnterprise = async (
   }
   const called = await execAsync(tcloudBin.value.bin, ['install_sqlmesh'], {
     signal: abortController.signal,
-    cwd: resolvedPath.value,
+    cwd: resolvedPath.value.workspaceFolder,
     env: tcloudBin.value.env,
   })
   if (called.exitCode !== 0) {
@@ -318,14 +318,14 @@ export const sqlmeshLspExec = async (): Promise<
       message: resolvedPath.error,
     })
   }
-  const workspacePath = resolvedPath.value
+  const workspacePath = resolvedPath.value.workspaceFolder
 
   const configuredLSPExec = getSqlmeshLspEntryPoint()
   if (configuredLSPExec) {
     traceLog(`Using configured SQLMesh LSP entry point: ${configuredLSPExec.entrypoint} ${configuredLSPExec.args.join(' ')}`)
     return ok({
       bin: configuredLSPExec.entrypoint,
-      workspacePath,
+      workspacePath: workspacePath,
       env: process.env,
       args: configuredLSPExec.args,
     })
@@ -381,7 +381,7 @@ export const sqlmeshLspExec = async (): Promise<
       if (isSemVerGreaterThanOrEqual(tcloudBinVersion.value, [2, 10, 1])) {
         return ok ({
           bin: tcloudBin.value.bin,
-          workspacePath,
+          workspacePath: workspacePath,
           env: tcloudBin.value.env,
           args: ['sqlmesh_lsp'],
         })
@@ -407,7 +407,7 @@ export const sqlmeshLspExec = async (): Promise<
     }
     return ok({
       bin: binPath,
-      workspacePath,
+      workspacePath: workspacePath,
       env: env.value,
       args: [],
     })
@@ -427,7 +427,7 @@ export const sqlmeshLspExec = async (): Promise<
     }
     return ok({
       bin: sqlmeshLSP,
-      workspacePath,
+      workspacePath: workspacePath,
       env: env.value,
       args: [],
     })

vscode/extension/tests/lineage.spec.ts

@@ -40,7 +40,7 @@ test('Lineage panel renders correctly - relative project path', async ({
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
 
   const settings = {
-    'sqlmesh.projectPath': './projects/sushi',
+    'sqlmesh.projectPaths': ['./projects/sushi'],
     'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
   await fs.ensureDir(path.join(tempDir, '.vscode'))
@@ -67,7 +67,7 @@ test('Lineage panel renders correctly - absolute project path', async ({
   await fs.copy(SUSHI_SOURCE_PATH, projectDir)
 
   const settings = {
-    'sqlmesh.projectPath': projectDir,
+    'sqlmesh.projectPaths': [projectDir],
     'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
   await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
@@ -90,7 +90,7 @@ test('Lineage panel renders correctly - relative project outside of workspace',
   await fs.ensureDir(workspaceDir)
 
   const settings = {
-    'sqlmesh.projectPath': './../projects/sushi',
+    'sqlmesh.projectPaths': ['./../projects/sushi'],
     'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
   await fs.ensureDir(path.join(workspaceDir, '.vscode'))
@@ -115,7 +115,7 @@ test('Lineage panel renders correctly - absolute path project outside of workspa
   await fs.ensureDir(workspaceDir)
 
   const settings = {
-    'sqlmesh.projectPath': projectDir,
+    'sqlmesh.projectPaths': [projectDir],
     'python.defaultInterpreterPath': sharedCodeServer.defaultPythonInterpreter,
   }
   await fs.ensureDir(path.join(workspaceDir, '.vscode'))

vscode/extension/tests/multi_project.spec.ts

@@ -0,0 +1,35 @@
+import { test } from './fixtures'
+import {
+  MULTI_SOURCE_PATH,
+  openServerPage,
+  waitForLoadedSQLMesh,
+} from './utils'
+import fs from 'fs-extra'
+
+test('should work with multi-project setups', async ({
+  page,
+  sharedCodeServer,
+  tempDir,
+}) => {
+  await fs.copy(MULTI_SOURCE_PATH, tempDir)
+
+  // Open the server
+  await openServerPage(page, tempDir, sharedCodeServer)
+
+  // Open a model
+  await page
+    .getByRole('treeitem', { name: 'repo_1', exact: true })
+    .locator('a')
+    .click()
+  await page
+    .getByRole('treeitem', { name: 'models', exact: true })
+    .locator('a')
+    .click()
+  await page
+    .getByRole('treeitem', { name: 'a.sql', exact: true })
+    .locator('a')
+    .click()
+
+  // Wait for for the project to be loaded
+  await waitForLoadedSQLMesh(page)
+})

vscode/extension/tests/utils.ts

@@ -14,6 +14,14 @@ export const SUSHI_SOURCE_PATH = path.join(
   'examples',
   'sushi',
 )
+export const MULTI_SOURCE_PATH = path.join(
+  __dirname,
+  '..',
+  '..',
+  '..',
+  'examples',
+  'multi',
+)
 export const REPO_ROOT = path.join(__dirname, '..', '..', '..')
 
 /**