Improve environment setup UX (#1395)

## Changes
- Let users proceed with dbconnect installation when there's mismatch in
dbr an python versions
- Show a warning message when there's mismatch in dbr and python
versions
- Focus configuration panel when users click on Databricks Connect
status bar button
- Show a notification when clicking on the above button when everything
is setup already


<img width="913" alt="Screenshot 2024-10-16 at 16 23 02"
src="https://github.com/user-attachments/assets/87ca7cbf-5a82-4b45-932f-7fb0e9ee4179">


## Tests
Manually and existing tests

Author: ilia-db

packages/databricks-vscode/src/feature-manager/FeatureManager.ts

@@ -20,6 +20,7 @@ export interface FeatureStepState {
     available: boolean;
     title?: string;
     message?: string;
+    warning?: string;
     action?: FeatureEnableAction;
     isDisabledByFf?: boolean;
 }

packages/databricks-vscode/src/feature-manager/MultiStepAccessVerfier.ts

@@ -76,12 +76,13 @@ export abstract class MultiStepAccessVerifier implements Feature {
         });
     }
 
-    acceptStep(id: string, title?: string, message?: string) {
+    acceptStep(id: string, title?: string, message?: string, warning?: string) {
         return this.updateStep({
             id: id,
             available: true,
             title,
             message,
+            warning,
         });
     }
 

packages/databricks-vscode/src/language/EnvironmentCommands.ts

@@ -1,4 +1,4 @@
-import {window} from "vscode";
+import {window, commands} from "vscode";
 import {FeatureManager} from "../feature-manager/FeatureManager";
 import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
 import {Cluster} from "../sdk-extensions";
@@ -12,6 +12,7 @@ export class EnvironmentCommands {
     ) {}
 
     async setup(stepId?: string) {
+        commands.executeCommand("configurationView.focus");
         await window.withProgress(
             {location: {viewId: "configurationView"}},
             () => this._setup(stepId)
@@ -24,6 +25,9 @@ export class EnvironmentCommands {
             true
         );
         if (state.available) {
+            window.showInformationMessage(
+                "Python environment and Databricks Connect are already set up."
+            );
             return true;
         }
         for (const [, step] of state.steps) {

packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts

@@ -161,38 +161,6 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
 
     async checkPythonEnvironment(): Promise<FeatureStepState> {
         const env = await this.pythonExtension.pythonEnvironment;
-        const dbrVersionParts =
-            this.connectionManager.cluster?.dbrVersion || [];
-        // DBR 13 and 14 require python 3.10
-        if (
-            (dbrVersionParts[0] === 13 || dbrVersionParts[0] === 14) &&
-            !this.matchEnvironmentVersion(env, 3, 10)
-        ) {
-            return this.rejectStep(
-                "checkPythonEnvironment",
-                "Activate an environment with Python 3.10",
-                `The python version should match DBR ${
-                    dbrVersionParts[0]
-                } requirements. ${this.printEnvironment(env)}`,
-                this.selectPythonInterpreter.bind(this)
-            );
-        }
-        // DBR 15 requires python 3.11
-        if (
-            dbrVersionParts[0] === 15 &&
-            !this.matchEnvironmentVersion(env, 3, 11)
-        ) {
-            return this.rejectStep(
-                "checkPythonEnvironment",
-                "Activate an environment with Python 3.11",
-                `The version should match DBR ${
-                    dbrVersionParts[0]
-                } requirements. ${this.printEnvironment(env)}`,
-                this.selectPythonInterpreter.bind(this)
-            );
-        }
-        // If we don't know DBR version (no cluster is connected or new version is released and the extension isn't updated yet),
-        // we still check that environment is active and has python >= 3.10
         const envVersionTooLow =
             env?.version && (env.version.major !== 3 || env.version.minor < 10);
         const noEnvironment = !env?.environment;
@@ -215,10 +183,30 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
                 this.selectPythonInterpreter.bind(this)
             );
         }
+        const dbrVersionParts =
+            this.connectionManager.cluster?.dbrVersion || [];
+        let warning;
+        if (
+            (dbrVersionParts[0] === 13 || dbrVersionParts[0] === 14) &&
+            !this.matchEnvironmentVersion(env, 3, 10)
+        ) {
+            warning = `Use python 3.10 to match DBR ${
+                dbrVersionParts[0]
+            } requirements. ${this.printEnvironment(env)}`;
+        }
+        if (
+            dbrVersionParts[0] === 15 &&
+            !this.matchEnvironmentVersion(env, 3, 11)
+        ) {
+            warning = `Use python 3.11 to match DBR ${
+                dbrVersionParts[0]
+            } requirements. ${this.printEnvironment(env)}`;
+        }
         return this.acceptStep(
             "checkPythonEnvironment",
             `Active Environment: ${env.environment.name}`,
-            env.executable.uri?.fsPath
+            env.executable.uri?.fsPath,
+            warning
         );
     }
 

packages/databricks-vscode/src/test/e2e/wdio.conf.ts

@@ -8,7 +8,7 @@ import path from "node:path";
 import {fileURLToPath} from "url";
 import assert from "assert";
 import fs from "fs/promises";
-import {ApiError, Config, WorkspaceClient} from "@databricks/databricks-sdk";
+import {Config, WorkspaceClient} from "@databricks/databricks-sdk";
 import * as ElementCustomCommands from "./customCommands/elementCustomCommands.ts";
 import {execFile as execFileCb} from "node:child_process";
 import {cpSync, mkdirSync, rmSync} from "node:fs";
@@ -588,26 +588,44 @@ function getWorkspaceClient(config: Config) {
 
 async function startCluster(
     workspaceClient: WorkspaceClient,
-    clusterId: string
+    clusterId: string,
+    attempt = 0
 ) {
-    console.log(`Starting cluster: ${clusterId}`);
-
-    try {
-        await (
-            await workspaceClient.clusters.start({
-                cluster_id: clusterId,
-            })
-        ).wait({
-            onProgress: async (state) => {
-                console.log(`Cluster state: ${state.state}`);
-            },
-        });
-    } catch (e: unknown) {
-        if (!(e instanceof ApiError && e.message.includes("INVALID_STATE"))) {
-            throw e;
-        }
-        console.log(e.message);
+    console.log(`Cluster ID: ${clusterId}`);
+    if (attempt > 100) {
+        throw new Error("Failed to start the cluster: too many attempts");
+    }
+    const cluster = await workspaceClient.clusters.get({
+        cluster_id: clusterId,
+    });
+    console.log(`Cluster State: ${cluster.state}`);
+    switch (cluster.state) {
+        case "RUNNING":
+            console.log("Cluster is already running");
+            break;
+        case "TERMINATED":
+        case "ERROR":
+        case "UNKNOWN":
+            console.log("Starting the cluster...");
+            await (
+                await workspaceClient.clusters.start({
+                    cluster_id: clusterId,
+                })
+            ).wait({
+                onProgress: async (state) => {
+                    console.log(`Cluster state: ${state.state}`);
+                },
+            });
+            break;
+        case "PENDING":
+        case "RESIZING":
+        case "TERMINATING":
+        case "RESTARTING":
+            console.log("Waiting and retrying...");
+            await sleep(10000);
+            await startCluster(workspaceClient, clusterId, attempt + 1);
+            break;
+        default:
+            throw new Error(`Unknown cluster state: ${cluster.state}`);
     }
-
-    console.log(`Cluster started`);
 }

packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts

@@ -64,7 +64,7 @@ export class EnvironmentComponent extends BaseComponent {
         const environmentState = await this.featureManager.isEnabled(
             "environment.dependencies"
         );
-        const children = [];
+        const children: ConfigurationTreeItem[] = [];
         for (const [id, step] of environmentState.steps) {
             if (!step.available) {
                 children.push({
@@ -93,6 +93,13 @@ export class EnvironmentComponent extends BaseComponent {
                     tooltip: step.message,
                     iconPath: new ThemeIcon("check"),
                 });
+                if (step.warning) {
+                    children.push({
+                        contextValue: getItemContext(id, true),
+                        label: step.warning,
+                        iconPath: new ThemeIcon("warning"),
+                    });
+                }
             }
         }
         return children;