Bundle selection ux (#1519)

## Changes When you land on a workspace without dabs config in the root, you can now select a sub-folder as the active bundle folder (without changing the workspace root, and without reloading the whole IDE). <img width="500" alt="Screenshot 2025-01-16 at 10 40 16" src="https://github.com/user-attachments/assets/5348c29a-1e8c-4ad3-b54e-bc3c59347922" /> The above is for the case when we can't find any sub-projects ourselves. Clicking on the top button opens an OS folder selection dialog. When we can detect sub projects, we show the "select sub-folder" button on top: <img width="500" alt="Screenshot 2025-01-16 at 10 39 44" src="https://github.com/user-attachments/assets/f96a0953-7e03-4281-b9a5-d80a1907bbe8" /> When you select the sub-folder and initialise the extension, we now show additional "Local Folder" UI at the top of the configuration: <img width="497" alt="Screenshot 2025-01-14 at 10 11 09" src="https://github.com/user-attachments/assets/17270619-c4f0-453f-9d9e-a85574ab3bd4" /> <img width="637" alt="Screenshot 2025-01-14 at 10 17 39" src="https://github.com/user-attachments/assets/af326541-8765-4583-a908-5fbd5cffd8b7" /> The UI for selecting sub-folders: <img width="880" alt="Screenshot 2025-01-14 at 10 11 52" src="https://github.com/user-attachments/assets/1eee2b12-8747-42da-ab72-89eb688ab4db" /> "Select another folder" opens a OS-level selection dialog. If you select a folder that's part of the workspace, we don't reload the IDE and just point our extension to it. When the selected folder is outside of the workspace, we reload the IDE with the folder being the new workspace root. After you select some sub-folder it's saved to workspace-scoped vscode storage, so the next time this workspace is opened our extension already knows what sub-folder to use. After the extension is initialised, you can now easily change the active bundle folder if you have multiple bundles in different sub-folders. ## Tests Manually and with a few new unit tests, e2e tests will be in the follow up PR --------- Co-authored-by: Julia Crawford (Databricks) <[email protected]>
databricks · Jan 16, 2025 · 6b5762e · 6b5762e
1 parent df60b3d
commit 6b5762e
Show file tree

Hide file tree

Showing 34 changed files with 519 additions and 409 deletions.
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -203,6 +203,13 @@
                 "enablement": "databricks.context.activated",
                 "category": "Databricks"
             },
+            {
+                "command": "databricks.bundle.selectActiveProjectFolder",
+                "icon": "$(gear)",
+                "title": "Select a Databricks project folder",
+                "enablement": "databricks.context.activated && databricks.context.bundle.deploymentState == idle",
+                "category": "Databricks"
+            },
             {
                 "command": "databricks.bundle.refreshRemoteState",
                 "icon": "$(refresh)",
@@ -444,17 +451,22 @@
             },
             {
                 "view": "configurationView",
-                "contents": "There are multiple Databricks projects in the folder:\n[Open existing Databricks Project](command:databricks.bundle.openSubProject)",
+                "contents": "Detected multiple Databricks projects in the VSCode workspace:\n[Select a project](command:databricks.bundle.selectActiveProjectFolder)",
                 "when": "workspaceFolderCount > 0 && databricks.context.initialized && databricks.context.subProjectsAvailable"
             },
             {
                 "view": "configurationView",
-                "contents": "Migrate current folder to a Databricks Project: \n[Migrate current folder to a Databricks Project](command:databricks.bundle.startManualMigration)",
+                "contents": "No Databricks project configuration detected in the root of the VSCode workspace:\n[Create configuration](command:databricks.bundle.startManualMigration)",
                 "when": "workspaceFolderCount > 0 && databricks.context.initialized && databricks.context.pendingManualMigration"
             },
             {
                 "view": "configurationView",
-                "contents": "[Create a new Databricks Project](command:databricks.bundle.initNewProject)",
+                "contents": "No Databricks projects detected in the VSCode workspace:\n[Select a project manually](command:databricks.bundle.selectActiveProjectFolder)",
+                "when": "workspaceFolderCount > 0 && databricks.context.initialized && !databricks.context.isBundleProject && !databricks.context.subProjectsAvailable"
+            },
+            {
+                "view": "configurationView",
+                "contents": "Chose a new parent folder and create a Databricks project based on a [template](https://docs.databricks.com/en/dev-tools/bundles/templates.html#databricks-asset-bundle-project-templates):\n[Create a new project](command:databricks.bundle.initNewProject)",
                 "when": "workspaceFolderCount > 0 && databricks.context.initialized && !databricks.context.isBundleProject"
             },
             {
@@ -464,7 +476,12 @@
             },
             {
                 "view": "configurationView",
-                "contents": "This folder is empty.\n[Create a new Databricks Project](command:databricks.bundle.initNewProject)",
+                "contents": "This folder is empty.\n[Open folder](command:vscode.openFolder)",
+                "when": "workspaceFolderCount == 0"
+            },
+            {
+                "view": "configurationView",
+                "contents": "[Create a new Databricks project](command:databricks.bundle.initNewProject)",
                 "when": "workspaceFolderCount == 0"
             },
             {
@@ -559,6 +576,16 @@
                     "when": "viewItem =~ /^databricks.*\\.(has-url).*$/ && databricks.context.bundle.deploymentState == idle",
                     "group": "navigation_2@0"
                 },
+                {
+                    "command": "databricks.bundle.selectActiveProjectFolder",
+                    "when": "viewItem =~ /^databricks.configuration.activeProjectFolder/ && databricks.context.bundle.deploymentState == idle",
+                    "group": "navigation_2@0"
+                },
+                {
+                    "command": "databricks.bundle.selectActiveProjectFolder",
+                    "when": "viewItem =~ /^databricks.configuration.activeProjectFolder/ && databricks.context.bundle.deploymentState == idle",
+                    "group": "inline@2"
+                },
                 {
                     "command": "databricks.connection.attachCluster",
                     "when": "view == clusterView && databricks.context.bundle.deploymentState == idle",

diff --git a/packages/databricks-vscode/resources/python/dbconnect-bootstrap.py b/packages/databricks-vscode/resources/python/dbconnect-bootstrap.py
@@ -3,23 +3,19 @@
 import logging
 from runpy import run_path
 
-# Load environment variables from .databricks/.databricks.env
-# We only look for the folder in the current working directory
-# since for all commands laucnhed from root workspace
-def load_env_file_from_cwd(path: str):
-    if not os.path.isdir(path):
-        return
-
-    env_file_path = os.path.join(path, ".databricks", ".databricks.env")
-    if not os.path.exists(os.path.dirname(env_file_path)):
-        return
-
-    with open(env_file_path, "r") as f:
-        for line in f.readlines():
-            key, value = line.strip().split("=", 1)
-            os.environ[key] = value
-    return
-
+def load_env_from_leaf(path: str) -> bool:
+    curdir = path if os.path.isdir(path) else os.path.dirname(path)
+    env_file_path = os.path.join(curdir, ".databricks", ".databricks.env")
+    if os.path.exists(env_file_path):
+        with open(env_file_path, "r") as f:
+            for line in f.readlines():
+                key, value = line.strip().split("=", 1)
+                os.environ[key] = value
+        return curdir
+    parent = os.path.dirname(curdir)
+    if parent == curdir:
+        return curdir
+    return load_env_from_leaf(parent)
 
 script = sys.argv[1]
 sys.argv = sys.argv[1:]
@@ -35,8 +31,7 @@ def load_env_file_from_cwd(path: str):
 # Suppress grpc warnings coming from databricks-connect with newer version of grpcio lib
 os.environ["GRPC_VERBOSITY"] = "NONE"
 
-root_dir = os.getcwd()
-load_env_file_from_cwd(root_dir)
+project_dir = load_env_from_leaf(cur_dir)
 
 log_level = os.environ.get("DATABRICKS_VSCODE_LOG_LEVEL")
 log_level = log_level if log_level is not None else "WARN"
@@ -68,7 +63,7 @@ def getArgument(*args, **kwargs):
 
 db_globals['getArgument'] = getArgument
 
-sys.path.insert(0, root_dir)
+sys.path.insert(0, project_dir)
 sys.path.insert(0, cur_dir)
 
 run_path(script, init_globals=db_globals, run_name="__main__")
diff --git a/packages/databricks-vscode/src/bundle/BundleFileSet.test.ts b/packages/databricks-vscode/src/bundle/BundleFileSet.test.ts
@@ -23,10 +23,12 @@ describe(__filename, async function () {
     function getWorkspaceFolderManagerMock() {
         const mockWorkspaceFolderManager = mock<WorkspaceFolderManager>();
         const mockWorkspaceFolder = mock<WorkspaceFolder>();
-        when(mockWorkspaceFolder.uri).thenReturn(Uri.file(tmpdir.path));
+        const uri = Uri.file(tmpdir.path);
+        when(mockWorkspaceFolder.uri).thenReturn(uri);
         when(mockWorkspaceFolderManager.activeWorkspaceFolder).thenReturn(
             instance(mockWorkspaceFolder)
         );
+        when(mockWorkspaceFolderManager.activeProjectUri).thenReturn(uri);
         return instance(mockWorkspaceFolderManager);
     }
 

diff --git a/packages/databricks-vscode/src/bundle/BundleFileSet.ts b/packages/databricks-vscode/src/bundle/BundleFileSet.ts
@@ -62,21 +62,21 @@ export class BundleFileSet {
             return bundle as BundleSchema;
         });
 
-    private get workspaceRoot() {
-        return this.workspaceFolderManager.activeWorkspaceFolder.uri;
+    private get projectRoot() {
+        return this.workspaceFolderManager.activeProjectUri;
     }
 
     constructor(
         private readonly workspaceFolderManager: WorkspaceFolderManager
     ) {
-        workspaceFolderManager.onDidChangeActiveWorkspaceFolder(() => {
+        workspaceFolderManager.onDidChangeActiveProjectFolder(() => {
             this.bundleDataCache.invalidate();
         });
     }
 
     async getRootFile() {
         const rootFile = await glob.glob(
-            getAbsoluteGlobPath(rootFilePattern, this.workspaceRoot),
+            getAbsoluteGlobPath(rootFilePattern, this.projectRoot),
             {nocase: process.platform === "win32"}
         );
         if (rootFile.length !== 1) {
@@ -85,33 +85,6 @@ export class BundleFileSet {
         return Uri.file(rootFile[0]);
     }
 
-    async getSubProjects(
-        root?: Uri
-    ): Promise<{relative: Uri; absolute: Uri}[]> {
-        const subProjectRoots = await glob.glob(
-            getAbsoluteGlobPath(
-                subProjectFilePattern,
-                root || this.workspaceRoot
-            ),
-            {nocase: process.platform === "win32"}
-        );
-        const normalizedRoot = path.normalize(
-            root?.fsPath ?? this.workspaceRoot.fsPath
-        );
-        return subProjectRoots
-            .map((rootFile) => {
-                const dirname = path.dirname(path.normalize(rootFile));
-                const absolute = Uri.file(dirname);
-                const relative = Uri.file(
-                    absolute.fsPath.replace(normalizedRoot, "")
-                );
-                return {absolute, relative};
-            })
-            .filter(({absolute}) => {
-                return absolute.fsPath !== normalizedRoot;
-            });
-    }
-
     async getIncludedFilesGlob() {
         const rootFile = await this.getRootFile();
         if (rootFile === undefined) {
@@ -133,7 +106,7 @@ export class BundleFileSet {
             return (
                 await glob.glob(
                     toGlobPath(
-                        path.join(this.workspaceRoot.fsPath, includedFilesGlob)
+                        path.join(this.projectRoot.fsPath, includedFilesGlob)
                     ),
                     {nocase: process.platform === "win32"}
                 )
@@ -171,7 +144,7 @@ export class BundleFileSet {
     isRootBundleFile(e: Uri) {
         return minimatch(
             e.fsPath,
-            getAbsoluteGlobPath(rootFilePattern, this.workspaceRoot)
+            getAbsoluteGlobPath(rootFilePattern, this.projectRoot)
         );
     }
 
@@ -182,7 +155,7 @@ export class BundleFileSet {
         }
         includedFilesGlob = getAbsoluteGlobPath(
             includedFilesGlob,
-            this.workspaceRoot
+            this.projectRoot
         );
         return minimatch(e.fsPath, toGlobPath(includedFilesGlob));
     }

diff --git a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts
@@ -15,48 +15,9 @@ import {getSubProjects} from "./BundleFileSet";
 import {tmpdir} from "os";
 import {ShellUtils} from "../utils";
 import {Events, Telemetry} from "../telemetry";
-import {OverrideableConfigModel} from "../configuration/models/OverrideableConfigModel";
-import {writeFile, mkdir} from "fs/promises";
-import path from "path";
 import {escapePathArgument} from "../utils/shellUtils";
-
-export async function promptToOpenSubProjects(
-    projects: {absolute: Uri; relative: string}[],
-    authProvider?: AuthProvider
-) {
-    type OpenProjectItem = QuickPickItem & {uri?: Uri};
-    const items: OpenProjectItem[] = projects.map((project) => {
-        return {
-            uri: project.absolute,
-            label: project.relative,
-            detail: project.absolute.fsPath,
-        };
-    });
-    items.push(
-        {label: "", kind: QuickPickItemKind.Separator},
-        {label: "Choose another folder"}
-    );
-    const options = {
-        title: "Select the project you want to open",
-    };
-    const item = await window.showQuickPick<OpenProjectItem>(items, options);
-    if (!item?.uri) {
-        return;
-    }
-
-    if (authProvider?.authType === "profile") {
-        const rootOverrideFilePath =
-            OverrideableConfigModel.getRootOverrideFile(item.uri);
-        await mkdir(path.dirname(rootOverrideFilePath.fsPath), {
-            recursive: true,
-        });
-        await writeFile(
-            rootOverrideFilePath.fsPath,
-            JSON.stringify({authProfile: authProvider.toJSON().profile})
-        );
-    }
-    await commands.executeCommand("vscode.openFolder", item.uri);
-}
+import {promptToSelectActiveProjectFolder} from "./activeBundleUtils";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 
 export class BundleInitWizard {
     private logger = logging.NamedLogger.getOrCreate(Loggers.Extension);
@@ -68,7 +29,8 @@ export class BundleInitWizard {
 
     public async initNewProject(
         workspaceUri?: Uri,
-        existingAuthProvider?: AuthProvider
+        existingAuthProvider?: AuthProvider,
+        workspaceFolderManager?: WorkspaceFolderManager
     ) {
         const recordEvent = this.telemetry.start(Events.BUNDLE_INIT);
         try {
@@ -97,7 +59,11 @@ export class BundleInitWizard {
                 this.logger.debug(
                     `Detected ${projects.length} sub projects after the init wizard, prompting to open one`
                 );
-                await promptToOpenSubProjects(projects, authProvider);
+                await promptToSelectActiveProjectFolder(
+                    projects,
+                    authProvider,
+                    workspaceFolderManager
+                );
             } else {
                 this.logger.debug(
                     `No projects detected after the init wizard, showing notification to open a folder manually`

diff --git a/packages/databricks-vscode/src/bundle/BundleProjectManager.ts b/packages/databricks-vscode/src/bundle/BundleProjectManager.ts
@@ -16,9 +16,10 @@ import {ProfileAuthProvider} from "../configuration/auth/AuthProvider";
 import {ProjectConfigFile} from "../file-managers/ProjectConfigFile";
 import {randomUUID} from "crypto";
 import {onError} from "../utils/onErrorDecorator";
-import {BundleInitWizard, promptToOpenSubProjects} from "./BundleInitWizard";
+import {BundleInitWizard} from "./BundleInitWizard";
 import {EventReporter, Events, Telemetry} from "../telemetry";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
+import {promptToSelectActiveProjectFolder} from "./activeBundleUtils";
 
 export class BundleProjectManager {
     private logger = logging.NamedLogger.getOrCreate(Loggers.Extension);
@@ -52,7 +53,7 @@ export class BundleProjectManager {
         private telemetry: Telemetry
     ) {
         this.disposables.push(
-            this.workspaceFolderManager.onDidChangeActiveWorkspaceFolder(
+            this.workspaceFolderManager.onDidChangeActiveProjectFolder(
                 async () => {
                     await this.isBundleProjectCache.refresh();
                 }
@@ -162,10 +163,18 @@ export class BundleProjectManager {
         });
     }
 
-    public async openSubProjects() {
-        if (this.subProjects && this.subProjects.length > 0) {
-            return promptToOpenSubProjects(this.subProjects);
-        }
+    public async selectActiveProjectFolder() {
+        return window.withProgress(
+            {location: {viewId: "configurationView"}},
+            async () => {
+                await this.detectSubProjects();
+                return promptToSelectActiveProjectFolder(
+                    this.subProjects ?? [],
+                    undefined,
+                    this.workspaceFolderManager
+                );
+            }
+        );
     }
 
     private setPendingManualMigration() {
@@ -328,7 +337,8 @@ export class BundleProjectManager {
             this.connectionManager.databricksWorkspace?.authProvider;
         const parentFolder = await bundleInitWizard.initNewProject(
             this.workspaceUri,
-            authProvider
+            authProvider,
+            this.workspaceFolderManager
         );
         if (parentFolder) {
             await this.isBundleProjectCache.refresh();

diff --git a/packages/databricks-vscode/src/bundle/BundleWatcher.ts b/packages/databricks-vscode/src/bundle/BundleWatcher.ts
@@ -25,7 +25,7 @@ export class BundleWatcher implements Disposable {
     ) {
         this.initCleanup = this.init();
         this.disposables.push(
-            this.workspaceFolderManager.onDidChangeActiveWorkspaceFolder(() => {
+            this.workspaceFolderManager.onDidChangeActiveProjectFolder(() => {
                 this.initCleanup.dispose();
                 this.initCleanup = this.init();
                 this.bundleFileSet.bundleDataCache.invalidate();
@@ -37,7 +37,7 @@ export class BundleWatcher implements Disposable {
         const yamlWatcher = workspace.createFileSystemWatcher(
             getAbsoluteGlobPath(
                 path.join("**", "*.{yaml,yml}"),
-                this.workspaceFolderManager.activeWorkspaceFolder.uri
+                this.workspaceFolderManager.activeProjectUri
             )
         );