Allow opting out of cluster override and make it more visible in UI (#1171)

## Changes



https://github.com/databricks/databricks-vscode/assets/88345179/2d2d4289-ebb3-4797-bffa-f01b209d8e4a


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

Author: kartikgupta-db

.github/workflows/tests.yml

@@ -77,7 +77,7 @@ jobs:
           working-directory: packages/databricks-vscode
 
       - name: Integration Tests
-        run: yarn run test:integ
+        run: yarn run test:integ:prepare && yarn run test:integ:run
         working-directory: packages/databricks-vscode
 
       - name: Integration Tests SDK wrappers

packages/databricks-vscode-types/package.json

@@ -23,7 +23,7 @@
         "test": "yarn run test:lint && yarn run test:unit"
     },
     "devDependencies": {
-        "@types/vscode": "^1.85.0",
+        "@types/vscode": "1.86.0",
         "eslint": "^8.55.0",
         "prettier": "^3.1.1",
         "typescript": "^5.3.3"

packages/databricks-vscode/package.json

@@ -841,17 +841,17 @@
         "@types/sinonjs__fake-timers": "^8.1.5",
         "@types/tmp": "^0.2.6",
         "@types/triple-beam": "^1.3.5",
-        "@types/vscode": "^1.83.0",
+        "@types/vscode": "1.86.0",
         "@types/yargs": "^17.0.32",
         "@typescript-eslint/eslint-plugin": "^6.14.0",
         "@typescript-eslint/parser": "^6.14.0",
         "@typescript-eslint/utils": "^6.14.0",
         "@vscode/test-electron": "^2.3.8",
-        "@wdio/cli": "^8.32.3",
-        "@wdio/local-runner": "^8.32.3",
-        "@wdio/mocha-framework": "^8.32.3",
-        "@wdio/spec-reporter": "^8.32.2",
-        "@wdio/types": "^8.32.2",
+        "@wdio/cli": "^8.35.1",
+        "@wdio/local-runner": "^8.35.1",
+        "@wdio/mocha-framework": "^8.35.0",
+        "@wdio/spec-reporter": "^8.32.4",
+        "@wdio/types": "^8.32.4",
         "chai": "^4.3.10",
         "esbuild": "^0.19.9",
         "eslint": "^8.55.0",

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

@@ -95,7 +95,9 @@ describe(__filename, async function () {
                     path.join(tmpdirUri.fsPath, "includes", "included.yaml")
                 ),
             ].map((v) => v.fsPath);
-            expect(actual).to.deep.equal(expected);
+            expect(Array.from(new Set(actual).values())).to.deep.equal(
+                Array.from(new Set(expected).values())
+            );
         });
 
         it("should return all bundle files", async () => {

packages/databricks-vscode/src/configuration/ConnectionCommands.ts

@@ -16,7 +16,10 @@ import {saveNewProfile} from "./LoginWizard";
 import {PersonalAccessTokenAuthProvider} from "./auth/AuthProvider";
 import {normalizeHost} from "../utils/urlUtils";
 import {CliWrapper, ProcessError} from "../cli/CliWrapper";
-import {AUTH_TYPE_SWITCH_ID, AUTH_TYPE_LOGIN_ID} from "./ui/AuthTypeComponent";
+import {
+    AUTH_TYPE_SWITCH_ID,
+    AUTH_TYPE_LOGIN_ID,
+} from "../ui/configuration-view/AuthTypeComponent";
 import {ManualLoginSource} from "../telemetry/constants";
 import {onError} from "../utils/onErrorDecorator";
 

packages/databricks-vscode/src/configuration/ConnectionManager.ts

@@ -119,7 +119,12 @@ export class ConnectionManager implements Disposable {
                       )
                     : undefined;
 
-            this.cli.setClusterId(clusterId);
+            if (
+                (await this.configModel.get("useClusterOverride")) ||
+                clusterId === undefined
+            ) {
+                this.cli.setClusterId(clusterId);
+            }
             this.onDidChangeClusterEmitter.fire(this.cluster);
         } catch (e) {
             this.configModel.set("clusterId", undefined);
@@ -144,6 +149,17 @@ export class ConnectionManager implements Disposable {
                     this.updateClusterManager,
                     this
                 ),
+                this.configModel.onDidChangeKey("useClusterOverride")(
+                    async () => {
+                        const useClusterOverride =
+                            await this.configModel.get("useClusterOverride");
+                        this.cli.setClusterId(
+                            useClusterOverride
+                                ? this._clusterManager?.cluster?.id
+                                : undefined
+                        );
+                    }
+                ),
                 // Don't just listen to target change for logging in. Also explictly listen for changes in the keys we care about.
                 // We don't have to listen to changes in authProfile as it's set by the login method and we don't respect other
                 // user changes.

packages/databricks-vscode/src/configuration/models/OverrideableConfigModel.ts

@@ -8,12 +8,13 @@ import {existsSync} from "fs";
 export type OverrideableConfigState = {
     authProfile?: string;
     clusterId?: string;
+    useClusterOverride?: boolean;
 };
 
 export function isOverrideableConfigKey(
     key: string
 ): key is keyof OverrideableConfigState {
-    return ["authProfile", "clusterId"].includes(key);
+    return ["authProfile", "clusterId", "useClusterOverride"].includes(key);
 }
 
 async function safeRead(filePath: Uri) {

packages/databricks-vscode/src/extension.ts

@@ -14,7 +14,7 @@ import {ConnectionManager} from "./configuration/ConnectionManager";
 import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider";
 import {ClusterModel} from "./cluster/ClusterModel";
 import {ClusterCommands} from "./cluster/ClusterCommands";
-import {ConfigurationDataProvider} from "./configuration/ui/ConfigurationDataProvider";
+import {ConfigurationDataProvider} from "./ui/configuration-view/ConfigurationDataProvider";
 import {RunCommands} from "./run/RunCommands";
 import {DatabricksDebugAdapterFactory} from "./run/DatabricksDebugAdapter";
 import {DatabricksWorkflowDebugAdapterFactory} from "./run/DatabricksWorkflowDebugAdapter";
@@ -64,10 +64,11 @@ import {BundleCommands} from "./ui/bundle-resource-explorer/BundleCommands";
 import {BundleRunTerminalManager} from "./bundle/run/BundleRunTerminalManager";
 import {BundleRunStatusManager} from "./bundle/run/BundleRunStatusManager";
 import {BundleProjectManager} from "./bundle/BundleProjectManager";
-import {TreeItemDecorationProvider} from "./ui/bundle-resource-explorer/DecorationProvider";
+import {TreeItemDecorationProvider} from "./ui/DecorationProvider";
 import {BundleInitWizard} from "./bundle/BundleInitWizard";
 import {DatabricksDebugConfigurationProvider} from "./run/DatabricksDebugConfigurationProvider";
 import {isIntegrationTest} from "./utils/developmentUtils";
+import {ConfigurationTreeViewManager} from "./ui/configuration-view/ConfigurationTreeViewManager";
 import {getCLIDependenciesEnvVars} from "./utils/envVarGenerators";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
@@ -482,6 +483,14 @@ export async function activate(
         configModel,
         cli
     );
+    const configurationView = window.createTreeView("configurationView", {
+        treeDataProvider: configurationDataProvider,
+    });
+
+    const configurationTreeViewManager = new ConfigurationTreeViewManager(
+        configurationView,
+        configModel
+    );
 
     const connectionCommands = new ConnectionCommands(
         workspaceFsCommands,
@@ -493,11 +502,8 @@ export async function activate(
 
     context.subscriptions.push(
         configurationDataProvider,
-
-        window.registerTreeDataProvider(
-            "configurationView",
-            configurationDataProvider
-        ),
+        configurationView,
+        configurationTreeViewManager,
         telemetry.registerCommand(
             "databricks.connection.bundle.selectTarget",
             connectionCommands.selectTarget,

packages/databricks-vscode/src/test/e2e/utils.ts

@@ -41,7 +41,11 @@ export async function findViewSection(name: ViewSectionType) {
     const views =
         (await (await control?.openView())?.getContent()?.getSections()) ?? [];
     for (const v of views) {
-        if ((await v.getTitle()).toUpperCase() === name) {
+        const title = await v.getTitle();
+        if (title === null) {
+            continue;
+        }
+        if (title.toUpperCase() === name) {
             return v;
         }
     }
@@ -50,17 +54,26 @@ export async function findViewSection(name: ViewSectionType) {
 export async function getViewSection(
     name: ViewSectionType
 ): Promise<ViewSection | undefined> {
-    const section = await findViewSection(name);
-    assert(section);
+    let section: ViewSection | undefined;
+    await browser.waitUntil(
+        async () => {
+            section = await findViewSection(name);
+            return section !== undefined;
+        },
+        {
+            timeout: 10 * 1000,
+            timeoutMsg: `Can't find view section "${name}"`,
+        }
+    );
 
     for (const s of ViewSectionTypes) {
         if (s !== name) {
             await (await findViewSection(s))?.collapse();
         }
     }
 
-    await section.expand();
-    await (await section.elem).click();
+    await section!.expand();
+    await (await section!.elem).click();
     return section;
 }
 

packages/databricks-vscode/src/ui/bundle-resource-explorer/DecorationProvider.ts renamed to packages/databricks-vscode/src/ui/DecorationProvider.ts

@@ -7,8 +7,8 @@ import {
     EventEmitter,
     Disposable,
 } from "vscode";
-import {BundleResourceExplorerTreeDataProvider} from "./BundleResourceExplorerTreeDataProvider";
-import {ConfigurationDataProvider} from "../../configuration/ui/ConfigurationDataProvider";
+import {BundleResourceExplorerTreeDataProvider} from "./bundle-resource-explorer/BundleResourceExplorerTreeDataProvider";
+import {ConfigurationDataProvider} from "./configuration-view/ConfigurationDataProvider";
 
 const SCHEME = "databricks-view-item";
 export class TreeItemDecorationProvider implements FileDecorationProvider {