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
databricks · Apr 22, 2024 · f00ebd2 · f00ebd2
1 parent b82c041
commit f00ebd2
Show file tree

Hide file tree

Showing 26 changed files with 470 additions and 193 deletions.
diff --git a/.github/workflows/tests.yml b/.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

diff --git a/packages/databricks-vscode-types/package.json b/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"

diff --git a/packages/databricks-vscode/package.json b/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",

diff --git a/packages/databricks-vscode/src/bundle/BundleFileSet.test.ts b/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 () => {

diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/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";
 

diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/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.

diff --git a/packages/databricks-vscode/src/configuration/models/OverrideableConfigModel.ts b/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) {

diff --git a/packages/databricks-vscode/src/extension.ts b/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,

diff --git a/packages/databricks-vscode/src/test/e2e/utils.ts b/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;
 }
 

diff --git a/...e-resource-explorer/DecorationProvider.ts → ...ricks-vscode/src/ui/DecorationProvider.ts b/...e-resource-explorer/DecorationProvider.ts → ...ricks-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 {

diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/JobTreeNode.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/JobTreeNode.ts
@@ -6,7 +6,8 @@ import {
     BundleResourceExplorerTreeNode,
 } from "./types";
 import {BundleRunStatusManager} from "../../bundle/run/BundleRunStatusManager";
-import {ContextUtils, DecorationUtils} from "./utils";
+import {ContextUtils} from "./utils";
+import {DecorationUtils} from "../utils";
 import {ConnectionManager} from "../../configuration/ConnectionManager";
 import {JobRunStatusTreeNode} from "./JobRunStatusTreeNode";
 import {JobRunStatus} from "../../bundle/run/JobRunStatus";

diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/PipelineTreeNode.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/PipelineTreeNode.ts
@@ -5,7 +5,9 @@ import {
     BundleResourceExplorerTreeNode,
 } from "./types";
 import {BundleRunStatusManager} from "../../bundle/run/BundleRunStatusManager";
-import {ContextUtils, DecorationUtils} from "./utils";
+import {ContextUtils} from "./utils";
+import {DecorationUtils} from "../utils";
+
 import {ConnectionManager} from "../../configuration/ConnectionManager";
 import {PipelineRunStatus} from "../../bundle/run/PipelineRunStatus";
 import {TreeItemTreeNode} from "./TreeItemTreeNode";

diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/LabelUtils.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/LabelUtils.ts
diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/index.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/index.ts
@@ -1,5 +1,3 @@
 export * as JobRunStateUtils from "./JobRunStateUtils";
 export * as RunStateUtils from "./RunStateUtils";
 export * as ContextUtils from "./ContextUtils";
-export * as LabelUtils from "./LabelUtils";
-export * as DecorationUtils from "./DecorationUtils";
diff --git a/...src/configuration/ui/AuthTypeComponent.ts → ...i/configuration-view/AuthTypeComponent.ts b/...src/configuration/ui/AuthTypeComponent.ts → ...i/configuration-view/AuthTypeComponent.ts
@@ -1,10 +1,11 @@
-import {ConfigModel} from "../models/ConfigModel";
-import {ConnectionManager} from "../ConnectionManager";
+import {ConfigModel} from "../../configuration/models/ConfigModel";
+import {ConnectionManager} from "../../configuration/ConnectionManager";
 import {BaseComponent} from "./BaseComponent";
 import {ConfigurationTreeItem} from "./types";
 import {ThemeIcon, ThemeColor} from "vscode";
-import {getProfilesForHost} from "../LoginWizard";
+import {getProfilesForHost} from "../../configuration/LoginWizard";
 import {CliWrapper} from "../../cli/CliWrapper";
+import {LabelUtils} from "../utils";
 
 export const AUTH_TYPE_SWITCH_ID = "AUTH-TYPE";
 export const AUTH_TYPE_LOGIN_ID = "LOGIN";
@@ -61,7 +62,7 @@ export class AuthTypeComponent extends BaseComponent {
             }
             return [
                 {
-                    label: {label, highlights: [[0, label.length]]},
+                    label: LabelUtils.highlightedLabel(label),
                     iconPath: new ThemeIcon(
                         "account",
                         new ThemeColor("notificationsErrorIcon.foreground")

diff --git a/...ode/src/configuration/ui/BaseComponent.ts → ...rc/ui/configuration-view/BaseComponent.ts b/...ode/src/configuration/ui/BaseComponent.ts → ...rc/ui/configuration-view/BaseComponent.ts
@@ -5,7 +5,6 @@ export abstract class BaseComponent implements Disposable {
     protected disposables: Disposable[] = [];
     protected onDidChangeEmitter = new EventEmitter<void>();
     public readonly onDidChange = this.onDidChangeEmitter.event;
-
     constructor() {}
 
     public abstract getChildren(

diff --git a/...configuration/ui/BundleTargetComponent.ts → ...nfiguration-view/BundleTargetComponent.ts b/...configuration/ui/BundleTargetComponent.ts → ...nfiguration-view/BundleTargetComponent.ts
@@ -1,8 +1,9 @@
 import {ThemeIcon, ThemeColor, TreeItemCollapsibleState, window} from "vscode";
-import {ConfigModel} from "../models/ConfigModel";
+import {ConfigModel} from "../../configuration/models/ConfigModel";
 import {BaseComponent} from "./BaseComponent";
 import {ConfigurationTreeItem} from "./types";
 import {UrlError} from "../../utils/urlUtils";
+import {LabelUtils} from "../utils";
 
 const TREE_ICON_ID = "TARGET";
 
@@ -30,13 +31,11 @@ export class BundleTargetComponent extends BaseComponent {
     private async getRoot(): Promise<ConfigurationTreeItem[]> {
         const target = this.configModel.target;
         if (target === undefined) {
-            const label = "Select a bundle target";
             return [
                 {
-                    label: {
-                        label,
-                        highlights: [[0, label.length]],
-                    },
+                    label: LabelUtils.highlightedLabel(
+                        "Select a bundle target"
+                    ),
                     id: TREE_ICON_ID,
                     iconPath: new ThemeIcon(
                         "plug",
@@ -82,10 +81,11 @@ export class BundleTargetComponent extends BaseComponent {
             ];
         } catch (e) {
             if (e instanceof UrlError) {
-                const label = `Invalid host for target ${target}`;
                 return [
                     {
-                        label: {label, highlights: [[0, label.length]]},
+                        label: LabelUtils.highlightedLabel(
+                            `Invalid host for target ${target}`
+                        ),
                         id: TREE_ICON_ID,
                         iconPath: new ThemeIcon(
                             "target",