Add experimental managed Python environment provisioning with uv

packages/databricks-vscode/package.json

@@ -1407,11 +1407,13 @@
                         "items": {
                             "enum": [
                                 "views.cluster",
-                                "views.workspace"
+                                "views.workspace",
+                                "python.managedEnvironment"
                             ],
                             "enumDescriptions": [
                                 "Show cluster view in the explorer.",
-                                "Show workspace browser in the explorer."
+                                "Show workspace browser in the explorer.",
+                                "Let the extension automatically provision a Python environment matching the selected compute when setting up Databricks Connect."
                             ],
                             "type": "string"
                         },

packages/databricks-vscode/src/extension.ts

@@ -75,6 +75,8 @@ import {BundleVariableTreeDataProvider} from "./ui/bundle-variables/BundleVariab
 import {ConfigurationTreeViewManager} from "./ui/configuration-view/ConfigurationTreeViewManager";
 import {getCLIDependenciesEnvVars} from "./utils/envVarGenerators";
 import {EnvironmentCommands} from "./language/EnvironmentCommands";
+import {EnvironmentProvisioner} from "./language/EnvironmentProvisioner";
+import {UvBinaryProvider} from "./language/UvBinaryProvider";
 import {WorkspaceFolderManager} from "./vscode-objs/WorkspaceFolderManager";
 import {SyncCommands} from "./sync/SyncCommands";
 import {CodeSynchronizer} from "./sync";
@@ -601,12 +603,21 @@ export async function activate(
                 configureAutocomplete
             )
     );
+    const environmentProvisioner = new EnvironmentProvisioner(
+        connectionManager,
+        pythonExtensionWrapper,
+        workspaceFolderManager,
+        new UvBinaryProvider(context),
+        telemetry
+    );
     const environmentCommands = new EnvironmentCommands(
         featureManager,
         pythonExtensionWrapper,
-        environmentDependenciesInstaller
+        environmentDependenciesInstaller,
+        environmentProvisioner
     );
     context.subscriptions.push(
+        environmentProvisioner,
         telemetry.registerCommand(
             "databricks.environment.setup",
             environmentCommands.setup,

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

@@ -0,0 +1,121 @@
+import * as assert from "assert";
+import {anything, instance, mock, verify, when} from "ts-mockito";
+import {EnvironmentCommands} from "./EnvironmentCommands";
+import {
+    FeatureManager,
+    FeatureState,
+    FeatureStepState,
+} from "../feature-manager/FeatureManager";
+import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
+import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
+import {EnvironmentProvisioner} from "./EnvironmentProvisioner";
+
+function makeState(steps: Array<Partial<FeatureStepState> & {id: string}>) {
+    const state: FeatureState = {
+        available: steps.every((s) => s.available),
+        steps: new Map(
+            steps.map((s) => [
+                s.id,
+                {available: false, ...s} as FeatureStepState,
+            ])
+        ),
+    };
+    return state;
+}
+
+describe(__filename, () => {
+    let featureManagerMock: FeatureManager;
+    let provisionerMock: EnvironmentProvisioner;
+    let commands: EnvironmentCommands;
+
+    beforeEach(() => {
+        featureManagerMock = mock<FeatureManager>(FeatureManager);
+        provisionerMock = mock(EnvironmentProvisioner);
+        commands = new EnvironmentCommands(
+            instance(featureManagerMock),
+            instance(mock(MsPythonExtensionWrapper)),
+            instance(mock(EnvironmentDependenciesInstaller)),
+            instance(provisionerMock)
+        );
+    });
+
+    describe("shouldProvision", () => {
+        it("should provision when only the python environment is failing", () => {
+            const state = makeState([
+                {id: "checkCluster", available: true},
+                {id: "checkPythonEnvironment", available: false},
+                {id: "checkEnvironmentDependencies", available: false},
+            ]);
+            assert.strictEqual(commands.shouldProvision(state), true);
+        });
+
+        it("should not provision when the cluster step is failing", () => {
+            const state = makeState([
+                {id: "checkCluster", available: false},
+                {id: "checkPythonEnvironment", available: false},
+            ]);
+            assert.strictEqual(commands.shouldProvision(state), false);
+        });
+
+        it("should not provision when everything is available", () => {
+            const state = makeState([
+                {id: "checkCluster", available: true},
+                {id: "checkPythonEnvironment", available: true},
+            ]);
+            assert.strictEqual(commands.shouldProvision(state), false);
+        });
+
+        it("should ignore failing optional steps", () => {
+            const state = makeState([
+                {id: "checkPythonEnvironment", available: false},
+                {id: "checkBuiltins", available: false, optional: true},
+            ]);
+            assert.strictEqual(commands.shouldProvision(state), true);
+        });
+
+        it("should respect the requested step", () => {
+            const state = makeState([
+                {id: "checkPythonEnvironment", available: false},
+            ]);
+            assert.strictEqual(
+                commands.shouldProvision(state, "checkPythonEnvironment"),
+                true
+            );
+            assert.strictEqual(
+                commands.shouldProvision(state, "checkCluster"),
+                false
+            );
+        });
+    });
+
+    describe("_setup", () => {
+        it("should provision when the experiment is enabled", async () => {
+            when(featureManagerMock.isEnabled(anything())).thenResolve(
+                makeState([{id: "checkPythonEnvironment", available: false}])
+            );
+            when(featureManagerMock.isEnabled(anything(), true)).thenResolve(
+                makeState([{id: "checkPythonEnvironment", available: true}])
+            );
+            when(provisionerMock.enabled).thenReturn(true);
+            when(provisionerMock.ensureEnvironment()).thenResolve({
+                success: true,
+            });
+
+            await commands["_setup"]();
+
+            verify(provisionerMock.ensureEnvironment()).once();
+            verify(featureManagerMock.isEnabled(anything(), true)).once();
+        });
+
+        it("should not provision when the experiment is disabled", async () => {
+            when(featureManagerMock.isEnabled(anything())).thenResolve(
+                makeState([{id: "checkPythonEnvironment", available: true}])
+            );
+            when(provisionerMock.enabled).thenReturn(false);
+
+            await commands["_setup"]();
+
+            verify(provisionerMock.ensureEnvironment()).never();
+        });
+    });
+});

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

@@ -9,12 +9,19 @@ import {Cluster} from "../sdk-extensions";
 import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller";
 import {Environment} from "./MsPythonExtensionApi";
 import {environmentName} from "../utils/environmentUtils";
+import {EnvironmentProvisioner} from "./EnvironmentProvisioner";
+
+const provisionableSteps = [
+    "checkPythonEnvironment",
+    "checkEnvironmentDependencies",
+];
 
 export class EnvironmentCommands {
     constructor(
         private featureManager: FeatureManager,
         private pythonExtension: MsPythonExtensionWrapper,
-        private installer: EnvironmentDependenciesInstaller
+        private installer: EnvironmentDependenciesInstaller,
+        private provisioner?: EnvironmentProvisioner
     ) {}
 
     async setup(stepId?: string): Promise<boolean> {
@@ -41,6 +48,20 @@ export class EnvironmentCommands {
         let state = await this.featureManager.isEnabled(
             "environment.dependencies"
         );
+        if (this.provisioner?.enabled && this.shouldProvision(state, stepId)) {
+            const result = await this.provisioner.ensureEnvironment();
+            if (!result.noOp) {
+                state = await this.checkEnvironmentDependencies();
+                if (state.available || !result.success) {
+                    // On provisioning failures the provisioner already showed
+                    // an actionable error with a retry option.
+                    this.reportSetupOutcome(state, !result.success);
+                    return state.available;
+                }
+            }
+            // noOp (or an inconsistent result): fall through to the manual
+            // per-step setup flow.
+        }
         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.
@@ -58,12 +79,29 @@ export class EnvironmentCommands {
         return state.available;
     }
 
-    private reportSetupOutcome(state: FeatureState) {
+    /**
+     * The managed flow can only fix the python environment and its
+     * dependencies: cluster and workspace problems keep the manual flow.
+     */
+    shouldProvision(state: FeatureState, stepId?: string): boolean {
+        if (stepId && !provisionableSteps.includes(stepId)) {
+            return false;
+        }
+        const failingSteps = Array.from(state.steps.values()).filter(
+            (s) => !s.available && !s.optional
+        );
+        return (
+            failingSteps.length > 0 &&
+            failingSteps.every((s) => provisionableSteps.includes(s.id))
+        );
+    }
+
+    private reportSetupOutcome(state: FeatureState, quiet = false) {
         if (state.available) {
             window.showInformationMessage(
                 "Python environment and Databricks Connect are set up."
             );
-        } else {
+        } else if (!quiet) {
             const detail = Array.from(state.steps.values())
                 .filter(
                     (s) => !s.available && !s.optional && (s.message || s.title)

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

@@ -5,6 +5,7 @@ import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {DATABRICKS_CONNECT_VERSION as DATABRICKS_CONNECT_MINIMAL_VERSION} from "../utils/constants";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
+import {resolveComputeTargetSpec} from "./computeTargetSpec";
 
 export class EnvironmentDependenciesInstaller implements Disposable {
     private disposables: Disposable[] = [];
@@ -67,23 +68,13 @@ export class EnvironmentDependenciesInstaller implements Disposable {
     }
 
     async getSuggestedVersion() {
-        if (this.connectionManager.serverless) {
-            const serverlessVersion =
-                workspaceConfigs.serverlessDbconnectVersion;
-            const parts = serverlessVersion.split(".");
-            const major = parts[0];
-            const minor = parts[1] ?? "3";
-            return `${major}.${minor}.*`;
-        }
-        const dbrVersionParts =
-            this.connectionManager.cluster?.dbrVersion || [];
-        if (dbrVersionParts.length < 2 || dbrVersionParts[0] === "x") {
-            return DATABRICKS_CONNECT_MINIMAL_VERSION;
-        }
-        const major = dbrVersionParts[0];
-        const minor = dbrVersionParts[1] === "x" ? "*" : dbrVersionParts[1];
-        const rest = minor === "*" ? "" : ".*";
-        return `${major}.${minor + rest}`;
+        const spec = resolveComputeTargetSpec({
+            serverless: this.connectionManager.serverless,
+            serverlessDbconnectVersion:
+                workspaceConfigs.serverlessDbconnectVersion,
+            dbrVersion: this.connectionManager.cluster?.dbrVersion,
+        });
+        return spec?.dbconnectVersion ?? DATABRICKS_CONNECT_MINIMAL_VERSION;
     }
 
     async installWithVersionPrompt(suggestedVersion?: string) {