Fix Python environment issues in the Databricks Connect run/debug flow

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

@@ -1,5 +1,9 @@
 import {window, commands, QuickPickItem, ProgressLocation} from "vscode";
-import {FeatureManager} from "../feature-manager/FeatureManager";
+import {
+    FeatureManager,
+    FeatureState,
+    FeatureStepState,
+} from "../feature-manager/FeatureManager";
 import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
 import {Cluster} from "../sdk-extensions";
 import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
@@ -13,9 +17,9 @@ export class EnvironmentCommands {
         private installer: EnvironmentDependenciesInstaller
     ) {}
 
-    async setup(stepId?: string) {
+    async setup(stepId?: string): Promise<boolean> {
         commands.executeCommand("configurationView.focus");
-        await window.withProgress(
+        return await window.withProgress(
             {location: {viewId: "configurationView"}},
             () => this._setup(stepId)
         );
@@ -32,7 +36,7 @@ export class EnvironmentCommands {
         );
     }
 
-    private async _setup(stepId?: string) {
+    private async _setup(stepId?: string): Promise<boolean> {
         // 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"
@@ -50,6 +54,11 @@ export class EnvironmentCommands {
                 break;
             }
         }
+        this.reportSetupOutcome(state);
+        return state.available;
+    }
+
+    private reportSetupOutcome(state: FeatureState) {
         if (state.available) {
             window.showInformationMessage(
                 "Python environment and Databricks Connect are set up."
@@ -77,14 +86,36 @@ export class EnvironmentCommands {
     async selectPythonInterpreter() {
         const environments =
             await this.pythonExtension.getAvailableEnvironments();
+        // The requirement hint is best effort: a fresh state check can block
+        // on the workspace connection, and the picker must always show up.
+        const state = await Promise.race([
+            this.featureManager.isEnabled("environment.dependencies"),
+            new Promise<undefined>((resolve) =>
+                setTimeout(() => resolve(undefined), 2000)
+            ),
+        ]);
+        const pythonStep = state?.steps.get("checkPythonEnvironment");
+        const requirement = !pythonStep?.available ? pythonStep : undefined;
         if (environments.length > 0) {
-            await this.showEnvironmentsQuickPick(environments);
+            await this.showEnvironmentsQuickPick(environments, requirement);
         } else {
-            await this.pythonExtension.createPythonEnvironment();
+            await this.createPythonEnvironment(requirement);
         }
     }
 
-    async showEnvironmentsQuickPick(environments: Environment[]) {
+    private async createPythonEnvironment(requirement?: FeatureStepState) {
+        if (requirement?.message) {
+            // The environment creation flow of the MS Python extension knows
+            // nothing about our version requirements, so we surface them here.
+            window.showInformationMessage(requirement.message);
+        }
+        await this.pythonExtension.createPythonEnvironment();
+    }
+
+    async showEnvironmentsQuickPick(
+        environments: Environment[],
+        requirement?: FeatureStepState
+    ) {
         const envPicks: (QuickPickItem & {path?: string})[] = environments.map(
             (env) => ({
                 label: environmentName(env),
@@ -101,11 +132,14 @@ export class EnvironmentCommands {
         ];
         const selectedPick = await window.showQuickPick(
             envPicks.concat(staticPicks),
-            {title: "Select Python Environment"}
+            {
+                title: requirement?.title ?? "Select Python Environment",
+                placeHolder: requirement?.message,
+            }
         );
         if (selectedPick) {
             if (selectedPick.label === createNewLabel) {
-                await this.pythonExtension.createPythonEnvironment();
+                await this.createPythonEnvironment(requirement);
             } else if (selectedPick.label === usePythonExtensionLabel) {
                 await this.pythonExtension.selectPythonInterpreter();
             } else if (selectedPick.path) {

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

@@ -0,0 +1,175 @@
+import * as assert from "assert";
+import {Disposable, Uri} from "vscode";
+import {mock, instance, when} from "ts-mockito";
+import {EnvironmentDependenciesVerifier} from "./EnvironmentDependenciesVerifier";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
+import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
+import {ConfigureAutocomplete} from "./ConfigureAutocomplete";
+import {ResolvedEnvironment} from "./MsPythonExtensionApi";
+import {Cluster} from "../sdk-extensions";
+
+function fakeEnvironment(
+    version?: {major: number; minor: number; micro: number},
+    name = ".venv",
+    executablePath = "/project/.venv/bin/python"
+): ResolvedEnvironment {
+    return {
+        id: executablePath,
+        path: executablePath,
+        executable: {uri: Uri.file(executablePath)},
+        environment: {
+            type: "VirtualEnvironment",
+            name,
+            folderUri: Uri.file("/project/.venv"),
+        },
+        version,
+    } as unknown as ResolvedEnvironment;
+}
+
+describe(__filename, () => {
+    let connectionManagerMock: ConnectionManager;
+    let pythonExtensionMock: MsPythonExtensionWrapper;
+    let verifier: EnvironmentDependenciesVerifier;
+
+    const noopEvent = () => new Disposable(() => {});
+
+    beforeEach(() => {
+        connectionManagerMock = mock(ConnectionManager);
+        when(connectionManagerMock.onDidChangeCluster).thenReturn(noopEvent);
+        when(connectionManagerMock.onDidChangeState).thenReturn(noopEvent);
+
+        pythonExtensionMock = mock(MsPythonExtensionWrapper);
+        when(pythonExtensionMock.onDidChangePythonExecutable).thenReturn(
+            noopEvent
+        );
+
+        const installerMock = mock(EnvironmentDependenciesInstaller);
+        when(installerMock.onDidTryInstallation).thenReturn(noopEvent);
+
+        const autocompleteMock = mock(ConfigureAutocomplete);
+        when(autocompleteMock.onDidUpdate).thenReturn(noopEvent);
+
+        verifier = new EnvironmentDependenciesVerifier(
+            instance(connectionManagerMock),
+            instance(pythonExtensionMock),
+            instance(installerMock),
+            instance(autocompleteMock)
+        );
+    });
+
+    function setupEnvironment(env?: ResolvedEnvironment) {
+        when(pythonExtensionMock.pythonEnvironment).thenReturn(
+            Promise.resolve(env)
+        );
+        when(pythonExtensionMock.getPythonExecutable()).thenResolve(
+            env?.executable.uri?.fsPath
+        );
+    }
+
+    describe("serverless (default dbconnect version 17.3)", () => {
+        beforeEach(() => {
+            when(connectionManagerMock.serverless).thenReturn(true);
+            when(connectionManagerMock.cluster).thenReturn(undefined);
+        });
+
+        it("should accept python 3.12 without warnings", async () => {
+            setupEnvironment(fakeEnvironment({major: 3, minor: 12, micro: 4}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, true);
+            assert.strictEqual(step.warning, undefined);
+            assert.strictEqual(step.title, "Active Environment: .venv");
+            assert.strictEqual(step.message, "/project/.venv/bin/python");
+        });
+
+        it("should reject python 3.13 (higher than the remote version)", async () => {
+            setupEnvironment(fakeEnvironment({major: 3, minor: 13, micro: 1}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, false);
+            assert.strictEqual(
+                step.title,
+                "Activate an environment with Python 3.12"
+            );
+            assert.ok(step.message?.includes("must match"));
+            assert.ok(
+                step.message?.includes("Current Python version is 3.13.1")
+            );
+        });
+
+        it("should reject python 3.9", async () => {
+            setupEnvironment(fakeEnvironment({major: 3, minor: 9, micro: 6}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, false);
+            assert.strictEqual(
+                step.title,
+                "Activate an environment with Python 3.12"
+            );
+        });
+
+        it("should reject when no environment is active", async () => {
+            setupEnvironment(undefined);
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, false);
+            assert.strictEqual(
+                step.title,
+                "Activate an environment with Python 3.12"
+            );
+            assert.ok(step.message?.includes("No active environments found"));
+        });
+
+        it("should accept an environment with an unresolvable version", async () => {
+            setupEnvironment(fakeEnvironment(undefined));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, true);
+            assert.strictEqual(step.warning, undefined);
+        });
+    });
+
+    describe("clusters", () => {
+        function setupCluster(dbrVersion: (number | "x")[]) {
+            when(connectionManagerMock.serverless).thenReturn(false);
+            when(connectionManagerMock.cluster).thenReturn({
+                dbrVersion,
+            } as unknown as Cluster);
+        }
+
+        it("should accept a matching python version", async () => {
+            setupCluster([15, 4]);
+            setupEnvironment(fakeEnvironment({major: 3, minor: 11, micro: 9}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, true);
+            assert.strictEqual(step.warning, undefined);
+        });
+
+        it("should accept a non-matching python >= 3.10 with a warning", async () => {
+            setupCluster([15, 4]);
+            setupEnvironment(fakeEnvironment({major: 3, minor: 10, micro: 2}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, true);
+            assert.ok(step.warning?.includes("Use Python 3.11"));
+            assert.ok(step.warning?.includes("DBR 15"));
+        });
+
+        it("should reject python below 3.10", async () => {
+            setupCluster([15, 4]);
+            setupEnvironment(fakeEnvironment({major: 3, minor: 9, micro: 6}));
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, false);
+            assert.strictEqual(
+                step.title,
+                "Activate an environment with Python 3.11"
+            );
+        });
+
+        it("should fall back to '3.10 or greater' for unknown DBR versions", async () => {
+            setupCluster(["x", "x"]);
+            setupEnvironment(undefined);
+            const step = await verifier.checkPythonEnvironment();
+            assert.strictEqual(step.available, false);
+            assert.strictEqual(
+                step.title,
+                "Activate an environment with Python 3.10 or greater"
+            );
+        });
+    });
+});