Add support for multiple vscode workspaces (#1250)

## Changes
<!-- Summary of your changes that are easy to understand -->

## Tests
<!-- How is this tested? -->

Author: kartikgupta-db

packages/databricks-vscode/.vscodeignore

@@ -18,3 +18,4 @@ logs/
 extension/
 **/*.vsix
 .build/
+tmp/**

packages/databricks-vscode/package.json

@@ -544,7 +544,7 @@
                 {
                     "submenu": "databricks.run",
                     "group": "navigation@0",
-                    "when": "resourceLangId == python || resourceLangId == scala || resourceLangId == r || resourceLangId == sql || resourceExtname == .ipynb"
+                    "when": "databricks.context.isActiveFileInActiveWorkspace && resourceLangId =~ /^(python|scala|r|sql)$/ || databricks.context.isActiveFileInActiveWorkspace && resourceExtname == .ipynb"
                 }
             ],
             "databricks.run": [
@@ -881,7 +881,7 @@
         "useYarn": false
     },
     "cli": {
-        "version": "0.219.0"
+        "version": "0.221.0"
     },
     "scripts": {
         "vscode:prepublish": "rm -rf out && yarn run package:compile && yarn run package:wrappers:write && yarn run package:jupyter-init-script:write && yarn run package:copy-webview-toolkit && yarn run generate-telemetry",

packages/databricks-vscode/src/bundle/BundleFileSet.test.ts

@@ -1,11 +1,13 @@
-import {Uri} from "vscode";
+import {Uri, WorkspaceFolder} from "vscode";
 import {BundleFileSet, getAbsoluteGlobPath} from "./BundleFileSet";
 import {expect} from "chai";
 import path from "path";
 import * as tmp from "tmp-promise";
 import * as fs from "fs/promises";
 import {BundleSchema} from "./types";
 import * as yaml from "yaml";
+import {instance, mock, when} from "ts-mockito";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 describe(__filename, async function () {
     let tmpdir: tmp.DirectoryResult;
@@ -18,6 +20,16 @@ describe(__filename, async function () {
         await tmpdir.cleanup();
     });
 
+    function getWorkspaceFolderManagerMock() {
+        const mockWorkspaceFolderManager = mock<WorkspaceFolderManager>();
+        const mockWorkspaceFolder = mock<WorkspaceFolder>();
+        when(mockWorkspaceFolder.uri).thenReturn(Uri.file(tmpdir.path));
+        when(mockWorkspaceFolderManager.activeWorkspaceFolder).thenReturn(
+            instance(mockWorkspaceFolder)
+        );
+        return instance(mockWorkspaceFolderManager);
+    }
+
     it("should return the correct absolute glob path", () => {
         const tmpdirUri = Uri.file(tmpdir.path);
         let expectedGlob = path.join(tmpdirUri.fsPath, "test.txt");
@@ -34,7 +46,9 @@ describe(__filename, async function () {
 
     it("should find the correct root bundle yaml", async () => {
         const tmpdirUri = Uri.file(tmpdir.path);
-        const bundleFileSet = new BundleFileSet(tmpdirUri);
+        const bundleFileSet = new BundleFileSet(
+            getWorkspaceFolderManagerMock()
+        );
 
         expect(await bundleFileSet.getRootFile()).to.be.undefined;
 
@@ -47,7 +61,9 @@ describe(__filename, async function () {
 
     it("should return undefined if more than one root bundle yaml is found", async () => {
         const tmpdirUri = Uri.file(tmpdir.path);
-        const bundleFileSet = new BundleFileSet(tmpdirUri);
+        const bundleFileSet = new BundleFileSet(
+            getWorkspaceFolderManagerMock()
+        );
 
         await fs.writeFile(path.join(tmpdirUri.fsPath, "bundle.yaml"), "");
         await fs.writeFile(path.join(tmpdirUri.fsPath, "databricks.yaml"), "");
@@ -80,7 +96,9 @@ describe(__filename, async function () {
 
         it("should return correct included files", async () => {
             const tmpdirUri = Uri.file(tmpdir.path);
-            const bundleFileSet = new BundleFileSet(tmpdirUri);
+            const bundleFileSet = new BundleFileSet(
+                getWorkspaceFolderManagerMock()
+            );
 
             expect(await bundleFileSet.getIncludedFilesGlob()).to.equal(
                 `{included.yaml,${path.join("includes", "**", "*.yaml")}}`
@@ -102,7 +120,9 @@ describe(__filename, async function () {
 
         it("should return all bundle files", async () => {
             const tmpdirUri = Uri.file(tmpdir.path);
-            const bundleFileSet = new BundleFileSet(tmpdirUri);
+            const bundleFileSet = new BundleFileSet(
+                getWorkspaceFolderManagerMock()
+            );
 
             const actual = (await bundleFileSet.allFiles()).map(
                 (v) => v.fsPath
@@ -117,7 +137,9 @@ describe(__filename, async function () {
 
         it("isRootBundleFile should return true only for root bundle file", async () => {
             const tmpdirUri = Uri.file(tmpdir.path);
-            const bundleFileSet = new BundleFileSet(tmpdirUri);
+            const bundleFileSet = new BundleFileSet(
+                getWorkspaceFolderManagerMock()
+            );
 
             const possibleRoots = [
                 "bundle.yaml",
@@ -143,7 +165,9 @@ describe(__filename, async function () {
 
         it("isIncludedBundleFile should return true only for included files", async () => {
             const tmpdirUri = Uri.file(tmpdir.path);
-            const bundleFileSet = new BundleFileSet(tmpdirUri);
+            const bundleFileSet = new BundleFileSet(
+                getWorkspaceFolderManagerMock()
+            );
 
             expect(
                 await bundleFileSet.isIncludedBundleFile(
@@ -168,7 +192,9 @@ describe(__filename, async function () {
 
         it("isBundleFile should return true only for bundle files", async () => {
             const tmpdirUri = Uri.file(tmpdir.path);
-            const bundleFileSet = new BundleFileSet(tmpdirUri);
+            const bundleFileSet = new BundleFileSet(
+                getWorkspaceFolderManagerMock()
+            );
 
             const possibleBundleFiles = [
                 "bundle.yaml",

packages/databricks-vscode/src/bundle/BundleFileSet.ts

@@ -7,6 +7,7 @@ import {BundleSchema} from "./types";
 import {readFile, writeFile} from "fs/promises";
 import {CachedValue} from "../locking/CachedValue";
 import minimatch from "minimatch";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 const rootFilePattern: string = "{bundle,databricks}.{yaml,yml}";
 const subProjectFilePattern: string = path.join("**", rootFilePattern);
@@ -61,7 +62,17 @@ export class BundleFileSet {
             return bundle as BundleSchema;
         });
 
-    constructor(private readonly workspaceRoot: Uri) {}
+    private get workspaceRoot() {
+        return this.workspaceFolderManager.activeWorkspaceFolder.uri;
+    }
+
+    constructor(
+        private readonly workspaceFolderManager: WorkspaceFolderManager
+    ) {
+        workspaceFolderManager.onDidChangeActiveWorkspaceFolder(() => {
+            this.bundleDataCache.invalidate();
+        });
+    }
 
     async getRootFile() {
         const rootFile = await glob.glob(

packages/databricks-vscode/src/bundle/BundleProjectManager.ts

@@ -18,6 +18,7 @@ import {randomUUID} from "crypto";
 import {onError} from "../utils/onErrorDecorator";
 import {BundleInitWizard, promptToOpenSubProjects} from "./BundleInitWizard";
 import {EventReporter, Events, Telemetry} from "../telemetry";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 export class BundleProjectManager {
     private logger = logging.NamedLogger.getOrCreate(Loggers.Extension);
@@ -36,17 +37,26 @@ export class BundleProjectManager {
     private subProjects?: {relative: string; absolute: Uri}[];
     private legacyProjectConfig?: ProjectConfigFile;
 
+    get workspaceUri() {
+        return this.workspaceFolderManager.activeWorkspaceFolder.uri;
+    }
+
     constructor(
         private context: ExtensionContext,
         private cli: CliWrapper,
         private customWhenContext: CustomWhenContext,
         private connectionManager: ConnectionManager,
         private configModel: ConfigModel,
         private bundleFileSet: BundleFileSet,
-        private workspaceUri: Uri,
+        private workspaceFolderManager: WorkspaceFolderManager,
         private telemetry: Telemetry
     ) {
         this.disposables.push(
+            this.workspaceFolderManager.onDidChangeActiveWorkspaceFolder(
+                async () => {
+                    await this.isBundleProjectCache.refresh();
+                }
+            ),
             this.bundleFileSet.bundleDataCache.onDidChange(async () => {
                 try {
                     await this.isBundleProjectCache.refresh();

packages/databricks-vscode/src/bundle/BundleWatcher.ts

@@ -1,7 +1,7 @@
 import {Disposable, EventEmitter, Uri, workspace} from "vscode";
 import {BundleFileSet, getAbsoluteGlobPath} from "./BundleFileSet";
-import {WithMutex} from "../locking";
 import path from "path";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 export class BundleWatcher implements Disposable {
     private disposables: Disposable[] = [];
@@ -18,15 +18,30 @@ export class BundleWatcher implements Disposable {
     private readonly _onDidDelete = new EventEmitter<Uri>();
     public readonly onDidDelete = this._onDidDelete.event;
 
-    private bundleFileSet: WithMutex<BundleFileSet>;
+    private initCleanup: Disposable;
+    constructor(
+        private readonly bundleFileSet: BundleFileSet,
+        private readonly workspaceFolderManager: WorkspaceFolderManager
+    ) {
+        this.initCleanup = this.init();
+        this.disposables.push(
+            this.workspaceFolderManager.onDidChangeActiveWorkspaceFolder(() => {
+                this.initCleanup.dispose();
+                this.initCleanup = this.init();
+                this.bundleFileSet.bundleDataCache.invalidate();
+            })
+        );
+    }
 
-    constructor(bundleFileSet: BundleFileSet, workspaceUri: Uri) {
-        this.bundleFileSet = new WithMutex(bundleFileSet);
+    private init() {
         const yamlWatcher = workspace.createFileSystemWatcher(
-            getAbsoluteGlobPath(path.join("**", "*.{yaml,yml}"), workspaceUri)
+            getAbsoluteGlobPath(
+                path.join("**", "*.{yaml,yml}"),
+                this.workspaceFolderManager.activeWorkspaceFolder.uri
+            )
         );
 
-        this.disposables.push(
+        const disposables: Disposable[] = [
             yamlWatcher,
             yamlWatcher.onDidCreate((e) => {
                 this.yamlFileChangeHandler(e, "CREATE");
@@ -36,22 +51,28 @@ export class BundleWatcher implements Disposable {
             }),
             yamlWatcher.onDidDelete((e) => {
                 this.yamlFileChangeHandler(e, "DELETE");
-            })
-        );
+            }),
+        ];
+
+        return {
+            dispose: () => {
+                disposables.forEach((i) => i.dispose());
+            },
+        };
     }
 
     private async yamlFileChangeHandler(
         e: Uri,
         type: "CREATE" | "CHANGE" | "DELETE"
     ) {
-        if (!(await this.bundleFileSet.value.isBundleFile(e))) {
+        if (!(await this.bundleFileSet.isBundleFile(e))) {
             return;
         }
 
-        this.bundleFileSet.value.bundleDataCache.invalidate();
+        this.bundleFileSet.bundleDataCache.invalidate();
         this._onDidChange.fire();
         // to provide additional granularity, we also fire an event when the root bundle file changes
-        if (this.bundleFileSet.value.isRootBundleFile(e)) {
+        if (this.bundleFileSet.isRootBundleFile(e)) {
             this._onDidChangeRootFile.fire();
         }
         switch (type) {
@@ -66,5 +87,6 @@ export class BundleWatcher implements Disposable {
 
     dispose() {
         this.disposables.forEach((i) => i.dispose());
+        this.initCleanup.dispose();
     }
 }

packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts

@@ -1,4 +1,3 @@
-import {Uri} from "vscode";
 import {CliWrapper} from "../../cli/CliWrapper";
 import {BaseModelWithStateCache} from "../../configuration/models/BaseModelWithStateCache";
 import {Mutex} from "../../locking";
@@ -9,6 +8,7 @@ import lodash from "lodash";
 import {WorkspaceConfigs} from "../../vscode-objs/WorkspaceConfigs";
 import {logging} from "@databricks/databricks-sdk";
 import {Loggers} from "../../logger";
+import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager";
 
 /* eslint-disable @typescript-eslint/naming-convention */
 export type BundleResourceModifiedStatus = "created" | "deleted" | "updated";
@@ -45,9 +45,13 @@ export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemote
     protected mutex = new Mutex();
     private logger = logging.NamedLogger.getOrCreate(Loggers.Bundle);
 
+    get workspaceFolder() {
+        return this.workspaceFolderManager.activeWorkspaceFolder.uri;
+    }
+
     constructor(
         private readonly cli: CliWrapper,
-        private readonly workspaceFolder: Uri,
+        private readonly workspaceFolderManager: WorkspaceFolderManager,
         private readonly workspaceConfigs: WorkspaceConfigs
     ) {
         super();

packages/databricks-vscode/src/bundle/models/BundleValidateModel.ts

@@ -1,4 +1,3 @@
-import {Uri} from "vscode";
 import {BundleWatcher} from "../BundleWatcher";
 import {AuthProvider} from "../../configuration/auth/AuthProvider";
 import {Mutex} from "../../locking";
@@ -10,6 +9,7 @@ import {BaseModelWithStateCache} from "../../configuration/models/BaseModelWithS
 import {withOnErrorHandler} from "../../utils/onErrorDecorator";
 import {logging} from "@databricks/databricks-sdk";
 import {Loggers} from "../../logger";
+import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager";
 
 export type BundleValidateState = {
     clusterId?: string;
@@ -25,7 +25,7 @@ export class BundleValidateModel extends BaseModelWithStateCache<BundleValidateS
     constructor(
         private readonly bundleWatcher: BundleWatcher,
         private readonly cli: CliWrapper,
-        private readonly workspaceFolder: Uri
+        private readonly workspaceFolderManager: WorkspaceFolderManager
     ) {
         super();
         this.disposables.push(
@@ -62,7 +62,11 @@ export class BundleValidateModel extends BaseModelWithStateCache<BundleValidateS
     }
 
     protected async readState(): Promise<BundleValidateState> {
-        if (this.target === undefined || this.authProvider === undefined) {
+        if (
+            !this.target ||
+            !this.authProvider ||
+            !this.workspaceFolderManager.activeWorkspaceFolder
+        ) {
             return {};
         }
 
@@ -71,7 +75,7 @@ export class BundleValidateModel extends BaseModelWithStateCache<BundleValidateS
                 await this.cli.bundleValidate(
                     this.target,
                     this.authProvider,
-                    this.workspaceFolder,
+                    this.workspaceFolderManager.activeWorkspaceFolder?.uri,
                     workspaceConfigs.databrickscfgLocation,
                     this.logger
                 )

packages/databricks-vscode/src/bundle/models/BundleVariableModel.ts

@@ -8,6 +8,7 @@ import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
 import {Loggers} from "../../logger";
 import {onError} from "../../utils/onErrorDecorator";
 import {BundleValidateModel} from "./BundleValidateModel";
+import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager";
 
 export type BundleVariable = Required<BundleSchema>["variables"][string] & {
     valueInTarget?: string;
@@ -23,10 +24,15 @@ export class BundleVariableModel extends BaseModelWithStateCache<BundleVariableM
     protected mutex: Mutex = new Mutex();
     private target: string | undefined;
     private overrideFileWatcher: FileSystemWatcher | undefined;
+
+    get workspaceRoot() {
+        return this.workspaceFolderManager.activeWorkspaceFolder.uri;
+    }
+
     constructor(
         private readonly configModel: ConfigModel,
         private readonly bundleValidateModel: BundleValidateModel,
-        private readonly workspaceRoot: Uri
+        private readonly workspaceFolderManager: WorkspaceFolderManager
     ) {
         super();
         this.disposables.push(