Skip to content

Commit 6852bc9

Browse files
committed
Add VPEX one-click Databricks Connect environment setup
*Why* Setting up a Databricks Connect Python environment that matches the selected compute target is a manual, error-prone flow. This adds a guided one-click path (the "VPEX" demo) that provisions a matched .venv and adopts it automatically, with real UI feedback. *What* - New `VpexEnvironmentSetup` flow: pre-flight confirmation, phase-aware progress notification, dedicated "VPEX Demo" output channel, and success/failure pop-ups. Runs the real CLI in the project folder: `databricks dbconnect init` then `dbconnect sync` (--serverless v4 --profile <detected>), parsing the CLI's `=== Phase N: name ===` output to narrate each step. Cancellable (SIGTERM/SIGKILL). - On success, auto-adopts `.venv/bin/python` via the MS Python extension API (refresh + updateActiveEnvironmentPath), degrading gracefully if absent. - New commands `databricks.environment.setupVpex` / `.showVpexVersions`, a `$(rocket)` status-bar button, and a node under the "Python Environment" tree section — all triggering the same flow and reflecting the matched state on success. - MsPythonExtensionWrapper: resolve `uv` from ~/.local/bin / ~/.cargo/bin / $XDG_BIN_HOME when it isn't on the extension-host PATH, so uv venvs no longer fall back to the (absent) native pip. Note: the `dbconnect` subcommand currently exists only in the dev CLI build, so CLI_PATH is hardcoded to it pending a bundled CLI that ships `dbconnect` (marked DEMO/POC in code). Not for merge. *Verification* - tsc --noEmit, eslint src, prettier: clean - Unit suite (VS Code test host): 237 passing, 0 failing - Ran the real `databricks dbconnect init` against the demo project: phases preflight→resolve→fetch→parse-python-version→plan→apply→ensure-python all ok; final `uv sync` (PyPI-dependent) blocked only by no network access in this environment. Co-authored-by: Isaac
1 parent a27514c commit 6852bc9

8 files changed

Lines changed: 786 additions & 9 deletions

File tree

packages/databricks-vscode/package.json

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,17 @@
440440
"enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode",
441441
"category": "Databricks"
442442
},
443+
{
444+
"command": "databricks.environment.setupVpex",
445+
"title": "Set up Python Environment (VPEX)",
446+
"icon": "$(rocket)",
447+
"category": "Databricks"
448+
},
449+
{
450+
"command": "databricks.environment.showVpexVersions",
451+
"title": "Show Matched Versions (VPEX)",
452+
"category": "Databricks"
453+
},
443454
{
444455
"command": "databricks.environment.selectPythonInterpreter",
445456
"title": "Change Python environment",

packages/databricks-vscode/src/extension.ts

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ import {Events, Metadata} from "./telemetry/constants";
4949
import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller";
5050
import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits";
5151
import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton";
52+
import {VpexEnvironmentSetup} from "./language/VpexEnvironmentSetup";
53+
import {VpexStatusBarButton} from "./language/VpexStatusBarButton";
5254
import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager";
5355
import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue";
5456
import {
@@ -646,6 +648,31 @@ export async function activate(
646648
featureManager
647649
);
648650

651+
// VPEX demo flow: a parallel environment-setup command that runs
652+
// `databricks dbconnect init/sync`, narrates phases, and auto-adopts the
653+
// resulting .venv.
654+
const vpexEnvironmentSetup = new VpexEnvironmentSetup(
655+
context,
656+
pythonExtensionWrapper
657+
);
658+
const vpexStatusBarButton = new VpexStatusBarButton(vpexEnvironmentSetup);
659+
context.subscriptions.push(
660+
vpexEnvironmentSetup,
661+
vpexStatusBarButton,
662+
telemetry.registerCommand(
663+
"databricks.environment.setupVpex",
664+
async () => {
665+
await vpexEnvironmentSetup.setup();
666+
vpexStatusBarButton.update();
667+
}
668+
),
669+
telemetry.registerCommand(
670+
"databricks.environment.showVpexVersions",
671+
vpexEnvironmentSetup.showVersions,
672+
vpexEnvironmentSetup
673+
)
674+
);
675+
649676
const databricksEnvFileManager = new DatabricksEnvFileManager(
650677
workspaceFolderManager,
651678
featureManager,
@@ -728,7 +755,8 @@ export async function activate(
728755
configModel,
729756
cli,
730757
featureManager,
731-
workspaceFolderManager
758+
workspaceFolderManager,
759+
vpexEnvironmentSetup
732760
);
733761
const configurationView = window.createTreeView("configurationView", {
734762
treeDataProvider: configurationDataProvider,

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

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import * as childProcess from "node:child_process";
1515
import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
1616
import {execFile} from "../cli/CliWrapper";
1717
import fs from "node:fs";
18+
import os from "node:os";
1819
import path from "node:path";
1920

2021
export class MsPythonExtensionWrapper implements Disposable {
@@ -120,13 +121,49 @@ export class MsPythonExtensionWrapper implements Disposable {
120121
});
121122
}
122123

124+
// Cached resolved path to the uv executable ("uv" if it's on PATH).
125+
private _uvCommand?: string;
126+
127+
/**
128+
* Resolve the uv executable. The extension host often inherits a minimal
129+
* PATH that does NOT include the locations the official uv installer uses
130+
* (~/.local/bin, ~/.cargo/bin). When that happens `execFile("uv", ...)`
131+
* throws ENOENT, which previously made isUsingUv() return false and sent
132+
* package operations down the native-pip path — fatal for uv venvs, which
133+
* have no pip ("No module named pip").
134+
*
135+
* So: try "uv" on PATH first, then fall back to the known install dirs.
136+
*/
137+
private async uvCommand(): Promise<string | undefined> {
138+
if (this._uvCommand) {
139+
return this._uvCommand;
140+
}
141+
const candidates = [
142+
"uv", // on PATH (covers Homebrew, system installs)
143+
path.join(os.homedir(), ".local", "bin", "uv"), // official installer
144+
path.join(os.homedir(), ".cargo", "bin", "uv"), // cargo install
145+
];
146+
if (process.env.XDG_BIN_HOME) {
147+
candidates.splice(1, 0, path.join(process.env.XDG_BIN_HOME, "uv"));
148+
}
149+
for (const candidate of candidates) {
150+
try {
151+
await execFile(candidate, ["--version"]);
152+
this._uvCommand = candidate;
153+
return candidate;
154+
} catch (error) {
155+
// try the next candidate
156+
}
157+
}
158+
return undefined;
159+
}
160+
123161
async isUsingUv() {
124-
try {
125-
await execFile("uv", ["--version"]);
126-
return fs.existsSync(path.join(this.projectRoot, "uv.lock"));
127-
} catch (error) {
162+
const uv = await this.uvCommand();
163+
if (!uv) {
128164
return false;
129165
}
166+
return fs.existsSync(path.join(this.projectRoot, "uv.lock"));
130167
}
131168

132169
private async getPipCommandAndArgs(
@@ -136,8 +173,11 @@ export class MsPythonExtensionWrapper implements Disposable {
136173
): Promise<{command: string; args: string[]}> {
137174
const isUv = await this.isUsingUv();
138175
if (isUv) {
176+
// isUsingUv() only returns true when uvCommand() resolved, so this
177+
// is guaranteed to be defined here.
178+
const uv = (await this.uvCommand())!;
139179
return {
140-
command: "uv",
180+
command: uv,
141181
args: ["pip", ...baseArgs, "--python", executable],
142182
};
143183
}
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import * as assert from "assert";
2+
import * as fs from "fs";
3+
import * as os from "os";
4+
import * as path from "path";
5+
import {ExtensionContext} from "vscode";
6+
import {instance, mock, when} from "ts-mockito";
7+
import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper";
8+
import {VpexEnvironmentSetup} from "./VpexEnvironmentSetup";
9+
10+
// Reach into the private methods we want to unit test without going through
11+
// the full VS Code UI flow.
12+
interface VpexInternals {
13+
detectTarget(projectDir: string): {
14+
serverless: boolean;
15+
authProfile: string;
16+
};
17+
friendlyPhase(stepLabel: string, name: string): string;
18+
}
19+
20+
describe(__filename, () => {
21+
let pythonExtensionMock: MsPythonExtensionWrapper;
22+
let setup: VpexEnvironmentSetup;
23+
let internals: VpexInternals;
24+
let tmpDir: string;
25+
26+
beforeEach(() => {
27+
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "vpex-test-"));
28+
pythonExtensionMock = mock(MsPythonExtensionWrapper);
29+
when(pythonExtensionMock.projectRoot).thenReturn(tmpDir);
30+
const context = {
31+
asAbsolutePath: (rel: string) => path.join(tmpDir, rel),
32+
} as unknown as ExtensionContext;
33+
setup = new VpexEnvironmentSetup(
34+
context,
35+
instance(pythonExtensionMock)
36+
);
37+
internals = setup as unknown as VpexInternals;
38+
});
39+
40+
afterEach(() => {
41+
setup.dispose();
42+
fs.rmSync(tmpDir, {recursive: true, force: true});
43+
});
44+
45+
function writeOverrides(contents: string) {
46+
const dir = path.join(tmpDir, ".databricks", "bundle", "dev");
47+
fs.mkdirSync(dir, {recursive: true});
48+
fs.writeFileSync(path.join(dir, "vscode.overrides.json"), contents);
49+
}
50+
51+
describe("detectTarget", () => {
52+
it("reads serverless and profile from the overrides file", () => {
53+
writeOverrides(
54+
JSON.stringify({authProfile: "prod", serverless: true})
55+
);
56+
const target = internals.detectTarget(tmpDir);
57+
assert.strictEqual(target.serverless, true);
58+
assert.strictEqual(target.authProfile, "prod");
59+
});
60+
61+
it("treats a non-serverless override as cluster", () => {
62+
writeOverrides(
63+
JSON.stringify({authProfile: "dev", serverless: false})
64+
);
65+
const target = internals.detectTarget(tmpDir);
66+
assert.strictEqual(target.serverless, false);
67+
});
68+
69+
it("falls back to serverless/dev when the file is missing", () => {
70+
const target = internals.detectTarget(tmpDir);
71+
assert.deepStrictEqual(target, {
72+
serverless: true,
73+
authProfile: "dev",
74+
});
75+
});
76+
77+
it("falls back gracefully on malformed JSON", () => {
78+
writeOverrides("{ not valid json");
79+
const target = internals.detectTarget(tmpDir);
80+
assert.deepStrictEqual(target, {
81+
serverless: true,
82+
authProfile: "dev",
83+
});
84+
});
85+
});
86+
87+
describe("friendlyPhase", () => {
88+
it("maps real CLI phase headers to narrated messages, prefixed with the step", () => {
89+
assert.strictEqual(
90+
internals.friendlyPhase("init", "preflight"),
91+
"dbconnect init: checking prerequisites…"
92+
);
93+
assert.strictEqual(
94+
internals.friendlyPhase("init", "resolve"),
95+
"dbconnect init: resolving target…"
96+
);
97+
assert.strictEqual(
98+
internals.friendlyPhase("init", "fetch"),
99+
"dbconnect init: fetching constraints…"
100+
);
101+
assert.strictEqual(
102+
internals.friendlyPhase("init", "parse-python-version"),
103+
"dbconnect init: reading Python version…"
104+
);
105+
assert.strictEqual(
106+
internals.friendlyPhase("sync", "plan"),
107+
"dbconnect sync: planning pyproject.toml changes…"
108+
);
109+
assert.strictEqual(
110+
internals.friendlyPhase("sync", "provision"),
111+
"dbconnect sync: provisioning .venv via uv…"
112+
);
113+
});
114+
115+
it("falls back to a generic label for unknown phases", () => {
116+
assert.strictEqual(
117+
internals.friendlyPhase("init", "something-new"),
118+
"dbconnect init: something-new"
119+
);
120+
});
121+
});
122+
});

0 commit comments

Comments
 (0)