[WIP] VPEX: one-click Databricks Connect environment setup

packages/databricks-vscode/package.json

@@ -440,6 +440,17 @@
                 "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode",
                 "category": "Databricks"
             },
+            {
+                "command": "databricks.environment.setupVpex",
+                "title": "Set up Python Environment (VPEX)",
+                "icon": "$(rocket)",
+                "category": "Databricks"
+            },
+            {
+                "command": "databricks.environment.showVpexVersions",
+                "title": "Show Matched Versions (VPEX)",
+                "category": "Databricks"
+            },
             {
                 "command": "databricks.environment.selectPythonInterpreter",
                 "title": "Change Python environment",

packages/databricks-vscode/src/extension.ts

@@ -49,6 +49,8 @@ import {Events, Metadata} from "./telemetry/constants";
 import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller";
 import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits";
 import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton";
+import {VpexEnvironmentSetup} from "./language/VpexEnvironmentSetup";
+import {VpexStatusBarButton} from "./language/VpexStatusBarButton";
 import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager";
 import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue";
 import {
@@ -646,6 +648,31 @@ export async function activate(
         featureManager
     );
 
+    // VPEX demo flow: a parallel environment-setup command that runs
+    // `databricks dbconnect init/sync`, narrates phases, and auto-adopts the
+    // resulting .venv.
+    const vpexEnvironmentSetup = new VpexEnvironmentSetup(
+        context,
+        pythonExtensionWrapper
+    );
+    const vpexStatusBarButton = new VpexStatusBarButton(vpexEnvironmentSetup);
+    context.subscriptions.push(
+        vpexEnvironmentSetup,
+        vpexStatusBarButton,
+        telemetry.registerCommand(
+            "databricks.environment.setupVpex",
+            async () => {
+                await vpexEnvironmentSetup.setup();
+                vpexStatusBarButton.update();
+            }
+        ),
+        telemetry.registerCommand(
+            "databricks.environment.showVpexVersions",
+            vpexEnvironmentSetup.showVersions,
+            vpexEnvironmentSetup
+        )
+    );
+
     const databricksEnvFileManager = new DatabricksEnvFileManager(
         workspaceFolderManager,
         featureManager,
@@ -728,7 +755,8 @@ export async function activate(
         configModel,
         cli,
         featureManager,
-        workspaceFolderManager
+        workspaceFolderManager,
+        vpexEnvironmentSetup
     );
     const configurationView = window.createTreeView("configurationView", {
         treeDataProvider: configurationDataProvider,

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

@@ -15,6 +15,7 @@ import * as childProcess from "node:child_process";
 import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
 import {execFile} from "../cli/CliWrapper";
 import fs from "node:fs";
+import os from "node:os";
 import path from "node:path";
 
 export class MsPythonExtensionWrapper implements Disposable {
@@ -120,13 +121,49 @@ export class MsPythonExtensionWrapper implements Disposable {
         });
     }
 
+    // Cached resolved path to the uv executable ("uv" if it's on PATH).
+    private _uvCommand?: string;
+
+    /**
+     * Resolve the uv executable. The extension host often inherits a minimal
+     * PATH that does NOT include the locations the official uv installer uses
+     * (~/.local/bin, ~/.cargo/bin). When that happens `execFile("uv", ...)`
+     * throws ENOENT, which previously made isUsingUv() return false and sent
+     * package operations down the native-pip path — fatal for uv venvs, which
+     * have no pip ("No module named pip").
+     *
+     * So: try "uv" on PATH first, then fall back to the known install dirs.
+     */
+    private async uvCommand(): Promise<string | undefined> {
+        if (this._uvCommand) {
+            return this._uvCommand;
+        }
+        const candidates = [
+            "uv", // on PATH (covers Homebrew, system installs)
+            path.join(os.homedir(), ".local", "bin", "uv"), // official installer
+            path.join(os.homedir(), ".cargo", "bin", "uv"), // cargo install
+        ];
+        if (process.env.XDG_BIN_HOME) {
+            candidates.splice(1, 0, path.join(process.env.XDG_BIN_HOME, "uv"));
+        }
+        for (const candidate of candidates) {
+            try {
+                await execFile(candidate, ["--version"]);
+                this._uvCommand = candidate;
+                return candidate;
+            } catch (error) {
+                // try the next candidate
+            }
+        }
+        return undefined;
+    }
+
     async isUsingUv() {
-        try {
-            await execFile("uv", ["--version"]);
-            return fs.existsSync(path.join(this.projectRoot, "uv.lock"));
-        } catch (error) {
+        const uv = await this.uvCommand();
+        if (!uv) {
             return false;
         }
+        return fs.existsSync(path.join(this.projectRoot, "uv.lock"));
     }
 
     private async getPipCommandAndArgs(
@@ -136,8 +173,11 @@ export class MsPythonExtensionWrapper implements Disposable {
     ): Promise<{command: string; args: string[]}> {
         const isUv = await this.isUsingUv();
         if (isUv) {
+            // isUsingUv() only returns true when uvCommand() resolved, so this
+            // is guaranteed to be defined here.
+            const uv = (await this.uvCommand())!;
             return {
-                command: "uv",
+                command: uv,
                 args: ["pip", ...baseArgs, "--python", executable],
             };
         }

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

@@ -0,0 +1,122 @@
+import * as assert from "assert";
+import * as fs from "fs";
+import * as os from "os";
+import * as path from "path";
+import {ExtensionContext} from "vscode";
+import {instance, mock, when} from "ts-mockito";
+import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
+import {VpexEnvironmentSetup} from "./VpexEnvironmentSetup";
+
+// Reach into the private methods we want to unit test without going through
+// the full VS Code UI flow.
+interface VpexInternals {
+    detectTarget(projectDir: string): {
+        serverless: boolean;
+        authProfile: string;
+    };
+    friendlyPhase(stepLabel: string, name: string): string;
+}
+
+describe(__filename, () => {
+    let pythonExtensionMock: MsPythonExtensionWrapper;
+    let setup: VpexEnvironmentSetup;
+    let internals: VpexInternals;
+    let tmpDir: string;
+
+    beforeEach(() => {
+        tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "vpex-test-"));
+        pythonExtensionMock = mock(MsPythonExtensionWrapper);
+        when(pythonExtensionMock.projectRoot).thenReturn(tmpDir);
+        const context = {
+            asAbsolutePath: (rel: string) => path.join(tmpDir, rel),
+        } as unknown as ExtensionContext;
+        setup = new VpexEnvironmentSetup(
+            context,
+            instance(pythonExtensionMock)
+        );
+        internals = setup as unknown as VpexInternals;
+    });
+
+    afterEach(() => {
+        setup.dispose();
+        fs.rmSync(tmpDir, {recursive: true, force: true});
+    });
+
+    function writeOverrides(contents: string) {
+        const dir = path.join(tmpDir, ".databricks", "bundle", "dev");
+        fs.mkdirSync(dir, {recursive: true});
+        fs.writeFileSync(path.join(dir, "vscode.overrides.json"), contents);
+    }
+
+    describe("detectTarget", () => {
+        it("reads serverless and profile from the overrides file", () => {
+            writeOverrides(
+                JSON.stringify({authProfile: "prod", serverless: true})
+            );
+            const target = internals.detectTarget(tmpDir);
+            assert.strictEqual(target.serverless, true);
+            assert.strictEqual(target.authProfile, "prod");
+        });
+
+        it("treats a non-serverless override as cluster", () => {
+            writeOverrides(
+                JSON.stringify({authProfile: "dev", serverless: false})
+            );
+            const target = internals.detectTarget(tmpDir);
+            assert.strictEqual(target.serverless, false);
+        });
+
+        it("falls back to serverless/dev when the file is missing", () => {
+            const target = internals.detectTarget(tmpDir);
+            assert.deepStrictEqual(target, {
+                serverless: true,
+                authProfile: "dev",
+            });
+        });
+
+        it("falls back gracefully on malformed JSON", () => {
+            writeOverrides("{ not valid json");
+            const target = internals.detectTarget(tmpDir);
+            assert.deepStrictEqual(target, {
+                serverless: true,
+                authProfile: "dev",
+            });
+        });
+    });
+
+    describe("friendlyPhase", () => {
+        it("maps real CLI phase headers to narrated messages, prefixed with the step", () => {
+            assert.strictEqual(
+                internals.friendlyPhase("init", "preflight"),
+                "dbconnect init: checking prerequisites…"
+            );
+            assert.strictEqual(
+                internals.friendlyPhase("init", "resolve"),
+                "dbconnect init: resolving target…"
+            );
+            assert.strictEqual(
+                internals.friendlyPhase("init", "fetch"),
+                "dbconnect init: fetching constraints…"
+            );
+            assert.strictEqual(
+                internals.friendlyPhase("init", "parse-python-version"),
+                "dbconnect init: reading Python version…"
+            );
+            assert.strictEqual(
+                internals.friendlyPhase("sync", "plan"),
+                "dbconnect sync: planning pyproject.toml changes…"
+            );
+            assert.strictEqual(
+                internals.friendlyPhase("sync", "provision"),
+                "dbconnect sync: provisioning .venv via uv…"
+            );
+        });
+
+        it("falls back to a generic label for unknown phases", () => {
+            assert.strictEqual(
+                internals.friendlyPhase("init", "something-new"),
+                "dbconnect init: something-new"
+            );
+        });
+    });
+});