Show error notifications after failed environment setup (#1608)

## Changes

Also change the setup logic slightly: we were re-checking the steps
before taking an action, now we are doing it after the action. See
comments in the code.

Fixes #1601

## Tests

Existing tests and manually

Author: ilia-db

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

@@ -1,4 +1,4 @@
-import {window, commands, QuickPickItem} from "vscode";
+import {window, commands, QuickPickItem, ProgressLocation} from "vscode";
 import {FeatureManager} from "../feature-manager/FeatureManager";
 import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
 import {Cluster} from "../sdk-extensions";
@@ -21,35 +21,56 @@ export class EnvironmentCommands {
         );
     }
 
+    private async checkEnvironmentDependencies() {
+        return await window.withProgress(
+            {
+                location: ProgressLocation.Notification,
+                title: `Databricks: checking python environment`,
+            },
+            () =>
+                this.featureManager.isEnabled("environment.dependencies", true)
+        );
+    }
+
     private async _setup(stepId?: string) {
-        const state = await this.featureManager.isEnabled(
-            "environment.dependencies",
-            true
+        // Get the state from the cache, we will re-check the state after taking an action (e.g. asking a user to select a venv or install dbconnect).
+        let state = await this.featureManager.isEnabled(
+            "environment.dependencies"
         );
-        for (const [, step] of state.steps) {
-            if (step.available || (stepId && step.id !== stepId)) {
-                continue;
-            }
-            if (step.action) {
-                return await step.action();
-            } else if (step.message) {
-                window.showErrorMessage(step.message);
-                return false;
+        for (const [, s] of state.steps) {
+            if (!s.available && (!stepId || s.id === stepId) && s.action) {
+                // Take an action of a failed step and re-check all steps state afterwards.
+                // Re-checking is important to get the up to date `state.available` value.
+                // It also fixes problems when Python extension doesn't notify us about the environment changes,
+                // and we end up being stuck with the outdated UI warnings.
+                await s.action();
+                state = await this.checkEnvironmentDependencies();
+                // All actions usually require user input (and can be cancelled), so here we stop after the first one
+                // and let users re-run the setup based on the (updated) UI state in the Python Environment panel.
+                break;
             }
         }
         if (state.available) {
             window.showInformationMessage(
                 "Python environment and Databricks Connect are set up."
             );
-            return true;
+        } else {
+            const detail = Array.from(state.steps.values())
+                .filter(
+                    (s) => !s.available && !s.optional && (s.message || s.title)
+                )
+                .map((s) => s.message || s.title)
+                .join("\n");
+            window.showErrorMessage(
+                `Failed to set up Python environment for Databricks Connect:\n${detail}`
+            );
         }
     }
 
     async refresh() {
         await window.withProgress(
             {location: {viewId: "configurationView"}},
-            () =>
-                this.featureManager.isEnabled("environment.dependencies", true)
+            () => this.checkEnvironmentDependencies()
         );
     }
 

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

@@ -169,7 +169,7 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
 
     private getCurrentPythonVersionMessage(env?: ResolvedEnvironment): string {
         return env?.version && env.environment
-            ? `Current version is ${env.version.major}.${env.version.minor}.${env.version.micro}.`
+            ? `Current Python version is ${env.version.major}.${env.version.minor}.${env.version.micro}.`
             : "No active environments found.";
     }
 
@@ -211,22 +211,26 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
 
     async checkPythonEnvironment(): Promise<FeatureStepState> {
         try {
-            const dbrVersionParts =
-                this.connectionManager.cluster?.dbrVersion || [];
-            const expectedPythonVersion =
-                this.getExpectedPythonVersionMessage(dbrVersionParts);
             const env = await this.pythonExtension.pythonEnvironment;
-            const envVersionTooLow =
+            let envVersionTooLow =
                 env?.version &&
                 (env.version.major !== 3 || env.version.minor < 10);
-            const noEnvironment = !env?.environment;
-            const currentPythonVersionMessage =
-                this.getCurrentPythonVersionMessage(env);
-            if (noEnvironment || envVersionTooLow) {
+            let dbrVersion = this.connectionManager.cluster?.dbrVersion || [];
+            if (this.connectionManager.serverless) {
+                dbrVersion = [15, 1];
+                envVersionTooLow =
+                    env?.version &&
+                    (env.version.major !== 3 || env.version.minor < 11);
+            }
+            const expectedPythonVersion =
+                this.getExpectedPythonVersionMessage(dbrVersion);
+            if (!env?.environment || envVersionTooLow) {
                 return this.rejectStep(
                     "checkPythonEnvironment",
                     `Activate an environment with Python ${expectedPythonVersion}`,
-                    `Databricks Connect requires ${expectedPythonVersion}. ${currentPythonVersionMessage}`,
+                    `Databricks Connect requires Python ${expectedPythonVersion}. ${this.getCurrentPythonVersionMessage(
+                        env
+                    )}`,
                     this.selectPythonInterpreter.bind(this)
                 );
             }
@@ -235,14 +239,14 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
                 return this.rejectStep(
                     "checkPythonEnvironment",
                     `Activate an environment with Python ${expectedPythonVersion}`,
-                    "No python executable found",
+                    "No Python executable found.",
                     this.selectPythonInterpreter.bind(this)
                 );
             }
             const warning = this.getVersionMismatchWarning(
-                dbrVersionParts[0],
+                dbrVersion[0],
                 env,
-                currentPythonVersionMessage
+                this.getCurrentPythonVersionMessage(env)
             );
             return this.acceptStep(
                 "checkPythonEnvironment",
@@ -332,7 +336,7 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier {
                 return this.rejectStep(
                     "checkEnvironmentDependencies",
                     "Install databricks-connect",
-                    "databricks-connect package is not installed in the current environment",
+                    "databricks-connect package is not installed in the current environment.",
                     async (advertisement = false) =>
                         this.installer.show(advertisement)
                 );