From 6852bc9312d9b8d4b3992b8b53c8248aec84c40e Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Tue, 23 Jun 2026 13:29:48 +0200 Subject: [PATCH] Add VPEX one-click Databricks Connect environment setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit *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 ), 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 --- packages/databricks-vscode/package.json | 11 + packages/databricks-vscode/src/extension.ts | 30 +- .../src/language/MsPythonExtensionWrapper.ts | 50 +- .../src/language/VpexEnvironmentSetup.test.ts | 122 +++++ .../src/language/VpexEnvironmentSetup.ts | 499 ++++++++++++++++++ .../src/language/VpexStatusBarButton.ts | 48 ++ .../ConfigurationDataProvider.ts | 7 +- .../EnvironmentComponent.ts | 28 +- 8 files changed, 786 insertions(+), 9 deletions(-) create mode 100644 packages/databricks-vscode/src/language/VpexEnvironmentSetup.test.ts create mode 100644 packages/databricks-vscode/src/language/VpexEnvironmentSetup.ts create mode 100644 packages/databricks-vscode/src/language/VpexStatusBarButton.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index cb6d6bbe0..79385db89 100644 --- a/packages/databricks-vscode/package.json +++ b/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", diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 69ca1fc94..771a0fa8c 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/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, diff --git a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts b/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts index b7fb0ff69..aad4a5cc3 100644 --- a/packages/databricks-vscode/src/language/MsPythonExtensionWrapper.ts +++ b/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 { + 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], }; } diff --git a/packages/databricks-vscode/src/language/VpexEnvironmentSetup.test.ts b/packages/databricks-vscode/src/language/VpexEnvironmentSetup.test.ts new file mode 100644 index 000000000..9be711475 --- /dev/null +++ b/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" + ); + }); + }); +}); diff --git a/packages/databricks-vscode/src/language/VpexEnvironmentSetup.ts b/packages/databricks-vscode/src/language/VpexEnvironmentSetup.ts new file mode 100644 index 000000000..d85852c73 --- /dev/null +++ b/packages/databricks-vscode/src/language/VpexEnvironmentSetup.ts @@ -0,0 +1,499 @@ +import { + window, + commands, + ProgressLocation, + Progress, + CancellationToken, + Uri, + OutputChannel, + Disposable, + ExtensionContext, + EventEmitter, +} from "vscode"; +import {spawn, ChildProcessWithoutNullStreams} from "child_process"; +import * as fs from "fs"; +import * as path from "path"; +import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper"; +import {NamedLogger} from "@databricks/sdk-experimental/dist/logging"; +import {Loggers} from "../logger"; + +// --------------------------------------------------------------------------- +// VPEX environment setup +// +// A self-contained parallel entry point that drives the real +// `databricks dbconnect` CLI from real UI: a pre-flight confirmation, a +// phase-aware progress notification, an output channel for raw logs, +// result/error pop-ups, and auto-adoption of the resulting `.venv` via the +// MS Python extension. +// +// It runs, in the open project folder: +// databricks dbconnect init --serverless --profile +// databricks dbconnect sync --serverless --profile +// +// This is a demo flow kept separate from the FeatureManager-driven +// `databricks.environment.setup` command. +// --------------------------------------------------------------------------- + +// DEMO/POC: the `dbconnect` subcommand only exists in the dev CLI build, not in +// the CLI currently bundled with the extension (v1.2.0). Point at the dev +// binary explicitly until `dbconnect` ships in the bundled CLI, at which point +// this should switch to CliWrapper.cliPath. +const CLI_PATH = "/Users/grigory.panov/work/cli/bin/databricks"; + +const OVERRIDES_REL = path.join( + ".databricks", + "bundle", + "dev", + "vscode.overrides.json" +); +const VENV_PYTHON_REL = path.join(".venv", "bin", "python"); // macOS only + +// Serverless version the demo provisions. The serverless env version isn't +// recorded in the project, so we pin it (matches dbconnect-init.sh's default). +const SERVERLESS_VERSION = "v4"; + +// Human-readable matched state shown after a successful run. +const MATCHED_TARGET = "serverless-v4"; +const MATCHED_PYTHON = "3.12"; +const MATCHED_DBCONNECT = "17.3"; + +const PHASE_RE = /===\s*Phase\s+(\S+):\s*(.+?)\s*===/; + +interface Target { + serverless: boolean; + authProfile: string; +} + +interface CommandResult { + code: number | null; + stdout: string; + stderr: string; + canceled: boolean; +} + +export class VpexEnvironmentSetup implements Disposable { + private outputChannel: OutputChannel; + private disposables: Disposable[] = []; + + /** + * Set after a successful setup so a status bar item or tree node can + * reflect the freshly matched environment. + */ + private _ready = false; + public get ready(): boolean { + return this._ready; + } + + // Fired whenever `ready` changes, so observers (status bar, tree) refresh. + private readonly onDidChangeStateEmitter = new EventEmitter(); + public readonly onDidChangeState = this.onDidChangeStateEmitter.event; + + constructor( + private readonly context: ExtensionContext, + private readonly pythonExtension: MsPythonExtensionWrapper + ) { + this.outputChannel = window.createOutputChannel("VPEX Demo"); + this.disposables.push(this.outputChannel, this.onDidChangeStateEmitter); + } + + private get projectDir(): string { + return this.pythonExtension.projectRoot; + } + + async setup() { + const projectDir = this.projectDir; + if (!projectDir) { + this.showFailure( + "No workspace folder is open. Open the project folder and try again." + ); + return; + } + + if (!fs.existsSync(CLI_PATH)) { + this.showFailure( + `Databricks CLI not found at ${CLI_PATH}. ` + + "The dbconnect subcommand requires the dev CLI build." + ); + return; + } + + // --- Pre-flight: detect target from the .databricks overrides ------ + const target = this.detectTarget(projectDir); + const targetLabel = target.serverless + ? `serverless (→ ${MATCHED_TARGET})` + : "cluster"; + + const confirm = await window.showInformationMessage( + `Detected target: ${targetLabel}, profile "${target.authProfile}".\n\nSet up ${MATCHED_TARGET} environment?`, + {modal: true}, + "Set up" + ); + if (confirm !== "Set up") { + // Cancel aborts quietly. + return; + } + + // --- Run the CLI with a phase-aware progress notification ---------- + // init: create pyproject.toml + provision .venv. + // sync: merge managed dependencies + re-provision. + const profile = target.authProfile; + const commonArgs = [ + "--serverless", + SERVERLESS_VERSION, + "--profile", + profile, + ]; + const steps: {label: string; args: string[]}[] = [ + {label: "init", args: ["dbconnect", "init", ...commonArgs]}, + {label: "sync", args: ["dbconnect", "sync", ...commonArgs]}, + ]; + + this.outputChannel.clear(); + + const result = await window.withProgress( + { + location: ProgressLocation.Notification, + title: "Databricks: setting up Python environment", + cancellable: true, + }, + async (progress, token) => { + let last: CommandResult = { + code: 0, + stdout: "", + stderr: "", + canceled: false, + }; + for (const step of steps) { + if (token.isCancellationRequested) { + return {...last, canceled: true}; + } + progress.report({ + message: `Running dbconnect ${step.label}…`, + }); + this.outputChannel.appendLine( + `$ ${CLI_PATH} ${step.args.join(" ")}` + ); + this.outputChannel.appendLine(` cwd: ${projectDir}`); + this.outputChannel.appendLine(""); + last = await this.runCli( + step.args, + projectDir, + step.label, + progress, + token + ); + this.outputChannel.appendLine(""); + // Stop the sequence on cancellation or failure. + if (last.canceled || last.code !== 0) { + return last; + } + } + return last; + } + ); + + if (result.canceled) { + window.showWarningMessage("VPEX: environment setup canceled."); + return; + } + + if (result.code !== 0) { + const tail = this.tailLines(result.stderr || result.stdout, 6); + this.showFailure( + `Setup failed (exit code ${result.code}).\n${tail}` + ); + return; + } + + // --- Success: auto-adopt the interpreter, then announce ------------ + await this.adoptInterpreter(projectDir); + this._ready = true; + this.onDidChangeStateEmitter.fire(); + + const choice = await window.showInformationMessage( + `Environment ready: ${MATCHED_TARGET} (Python ${MATCHED_PYTHON}, databricks-connect ${MATCHED_DBCONNECT}).`, + "Select Interpreter", + "Show Logs" + ); + if (choice === "Select Interpreter") { + await commands.executeCommand( + "databricks.environment.selectPythonInterpreter" + ); + } else if (choice === "Show Logs") { + this.outputChannel.show(true); + } + } + + async showVersions() { + const projectDir = this.projectDir; + if (!projectDir) { + window.showErrorMessage("VPEX: no workspace folder is open."); + return; + } + const venvPython = path.join(projectDir, VENV_PYTHON_REL); + if (!fs.existsSync(venvPython)) { + window.showErrorMessage( + "VPEX: no .venv found yet — run 'Set up Python Environment' first." + ); + return; + } + + const code = + "import sys, databricks.connect as c; " + + "print('Python ' + sys.version.split()[0]); " + + "print('databricks-connect ' + getattr(c, '__version__', 'unknown'))"; + + const result = await new Promise<{out: string; code: number | null}>( + (resolve) => { + const child = spawn(venvPython, ["-c", code], { + cwd: projectDir, + }); + let out = ""; + child.stdout.on("data", (b: Buffer) => (out += b.toString())); + child.stderr.on("data", (b: Buffer) => (out += b.toString())); + child.on("close", (c2) => resolve({out, code: c2})); + child.on("error", (e) => resolve({out: e.message, code: 1})); + } + ); + + if (result.code !== 0) { + window.showErrorMessage( + `VPEX: could not read versions: ${result.out.trim()}` + ); + return; + } + window.showInformationMessage( + `Matched environment:\n${result.out.trim()}`, + {modal: true} + ); + } + + // ----------------------------------------------------------------------- + // CLI runner + // ----------------------------------------------------------------------- + + private runCli( + args: string[], + cwd: string, + stepLabel: string, + progress: Progress<{message?: string; increment?: number}>, + token: CancellationToken + ): Promise { + return new Promise((resolve) => { + let child: ChildProcessWithoutNullStreams; + try { + child = spawn(CLI_PATH, args, {cwd, env: process.env}); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + resolve({code: 1, stdout: "", stderr: msg, canceled: false}); + return; + } + + let stdout = ""; + let stderr = ""; + let canceled = false; + + const cancelSub = token.onCancellationRequested(() => { + canceled = true; + this.outputChannel.appendLine( + "\n[VPEX] Cancellation requested — killing process…" + ); + child.kill("SIGTERM"); + setTimeout(() => { + if (!child.killed) { + child.kill("SIGKILL"); + } + }, 2000); + }); + + // Phase lines can arrive split across chunks; buffer by line. + let lineBuf = ""; + const handlePhaseText = (text: string) => { + lineBuf += text; + let idx: number; + while ((idx = lineBuf.indexOf("\n")) >= 0) { + const line = lineBuf.slice(0, idx); + lineBuf = lineBuf.slice(idx + 1); + // Strip ANSI color codes the CLI emits. + // eslint-disable-next-line no-control-regex + const clean = line.replace(/\x1b\[[0-9;]*m/g, ""); + const m = PHASE_RE.exec(clean); + if (m) { + progress.report({ + message: this.friendlyPhase(stepLabel, m[2]), + }); + } + } + }; + + child.stdout.on("data", (buf: Buffer) => { + const text = buf.toString(); + stdout += text; + this.outputChannel.append(text); + handlePhaseText(text); + }); + + child.stderr.on("data", (buf: Buffer) => { + const text = buf.toString(); + stderr += text; + this.outputChannel.append(text); + }); + + child.on("error", (err) => { + cancelSub.dispose(); + resolve({ + code: 1, + stdout, + stderr: stderr + `\n[spawn error] ${err.message}`, + canceled, + }); + }); + + child.on("close", (code) => { + cancelSub.dispose(); + resolve({code, stdout, stderr, canceled}); + }); + }); + } + + // Map raw `=== Phase N: ===` headers from `databricks dbconnect` + // to narrated progress messages, prefixed with the current step. + private friendlyPhase(stepLabel: string, name: string): string { + const prefix = `dbconnect ${stepLabel}`; + const lower = name.toLowerCase(); + if (lower.includes("preflight")) { + return `${prefix}: checking prerequisites…`; + } + if (lower.includes("resolve")) { + return `${prefix}: resolving target…`; + } + if (lower.includes("fetch")) { + return `${prefix}: fetching constraints…`; + } + if (lower.includes("parse-python-version")) { + return `${prefix}: reading Python version…`; + } + if (lower.includes("plan")) { + return `${prefix}: planning pyproject.toml changes…`; + } + if (lower.includes("apply")) { + return `${prefix}: applying pyproject.toml changes…`; + } + if (lower.includes("ensure-python")) { + return `${prefix}: installing Python via uv…`; + } + if (lower.includes("provision")) { + return `${prefix}: provisioning .venv via uv…`; + } + if (lower.includes("validate")) { + return `${prefix}: validating environment…`; + } + return `${prefix}: ${name}`; + } + + // ----------------------------------------------------------------------- + // Target detection + // ----------------------------------------------------------------------- + + private detectTarget(projectDir: string): Target { + const fallback: Target = {serverless: true, authProfile: "dev"}; + try { + const raw = fs.readFileSync( + path.join(projectDir, OVERRIDES_REL), + "utf8" + ); + const json = JSON.parse(raw); + return { + serverless: json.serverless === true, + authProfile: + typeof json.authProfile === "string" + ? json.authProfile + : "dev", + }; + } catch (e) { + this.outputChannel.appendLine( + `[VPEX] Could not read ${OVERRIDES_REL}; assuming serverless/dev.` + ); + return fallback; + } + } + + // ----------------------------------------------------------------------- + // Interpreter adoption via the MS Python extension API + // ----------------------------------------------------------------------- + + private async adoptInterpreter(projectDir: string): Promise { + const venvPython = path.join(projectDir, VENV_PYTHON_REL); + if (!fs.existsSync(venvPython)) { + this.outputChannel.appendLine( + `[VPEX] Expected interpreter not found at ${venvPython}; skipping auto-select.` + ); + return; + } + + const environments = this.pythonExtension.api?.environments; + if (!environments) { + this.outputChannel.appendLine( + "[VPEX] Python extension API unavailable; skipping auto-select. " + + "Pick the interpreter manually via 'Change Python environment'." + ); + window.showWarningMessage( + "VPEX: Python extension not available — please select the .venv interpreter manually." + ); + return; + } + + try { + this.outputChannel.appendLine( + "[VPEX] Refreshing Python environments…" + ); + await environments.refreshEnvironments(); + + this.outputChannel.appendLine( + `[VPEX] Selecting interpreter: ${venvPython}` + ); + await environments.updateActiveEnvironmentPath( + venvPython, + Uri.file(projectDir) + ); + this.outputChannel.appendLine("[VPEX] Interpreter adopted."); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + this.outputChannel.appendLine(`[VPEX] Auto-select failed: ${msg}`); + NamedLogger.getOrCreate(Loggers.Extension).error( + "VPEX interpreter auto-select failed", + err + ); + window.showWarningMessage( + "VPEX: could not auto-select the interpreter — please pick .venv manually." + ); + } + } + + // ----------------------------------------------------------------------- + // Helpers + // ----------------------------------------------------------------------- + + private showFailure(message: string): void { + window + .showErrorMessage(`VPEX: ${message}`, "Show Logs") + .then((choice) => { + if (choice === "Show Logs") { + this.outputChannel.show(true); + } + }); + } + + private tailLines(text: string, n: number): string { + const lines = text + // eslint-disable-next-line no-control-regex + .replace(/\x1b\[[0-9;]*m/g, "") + .split("\n") + .filter((l) => l.trim().length > 0); + return lines.slice(-n).join("\n"); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } +} diff --git a/packages/databricks-vscode/src/language/VpexStatusBarButton.ts b/packages/databricks-vscode/src/language/VpexStatusBarButton.ts new file mode 100644 index 000000000..bf84daa15 --- /dev/null +++ b/packages/databricks-vscode/src/language/VpexStatusBarButton.ts @@ -0,0 +1,48 @@ +import {Disposable, StatusBarAlignment, StatusBarItem, window} from "vscode"; +import {VpexEnvironmentSetup} from "./VpexEnvironmentSetup"; + +// Persistent status bar button that triggers the VPEX environment setup flow. +// After a successful setup it reflects the matched state. +export class VpexStatusBarButton implements Disposable { + private disposables: Disposable[] = []; + private statusBarButton: StatusBarItem; + + constructor(private readonly setup: VpexEnvironmentSetup) { + this.statusBarButton = window.createStatusBarItem( + StatusBarAlignment.Left, + 999 + ); + this.statusBarButton.command = "databricks.environment.setupVpex"; + this.disposables.push(this.statusBarButton); + this.setIdle(); + this.statusBarButton.show(); + } + + private setIdle() { + this.statusBarButton.name = "Databricks Env (VPEX)"; + this.statusBarButton.text = "$(rocket) Databricks Env"; + this.statusBarButton.tooltip = + "Set up the Databricks Connect Python environment"; + } + + private setReady() { + this.statusBarButton.name = "Databricks Env (VPEX)"; + this.statusBarButton.text = + "$(check) Databricks Env: serverless-v4 · py3.12"; + this.statusBarButton.tooltip = + "Environment ready: serverless-v4 (Python 3.12, databricks-connect 17.3). Click to re-run setup."; + } + + // Refresh the button label from the setup flow's current state. + public update() { + if (this.setup.ready) { + this.setReady(); + } else { + this.setIdle(); + } + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } +} diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index 2fd63c23c..5f83ed4a0 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -20,6 +20,7 @@ import {logging} from "@databricks/sdk-experimental"; import {Loggers} from "../../logger"; import {FeatureManager} from "../../feature-manager/FeatureManager"; import {EnvironmentComponent} from "./EnvironmentComponent"; +import {VpexEnvironmentSetup} from "../../language/VpexEnvironmentSetup"; import {WorkspaceFolderComponent} from "./WorkspaceFolderComponent"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; import {CodeSynchronizer} from "../../sync"; @@ -50,7 +51,8 @@ export class ConfigurationDataProvider private readonly configModel: ConfigModel, private readonly cli: CliWrapper, private readonly featureManager: FeatureManager, - private readonly workspaceFolderManager: WorkspaceFolderManager + private readonly workspaceFolderManager: WorkspaceFolderManager, + private readonly vpexEnvironmentSetup: VpexEnvironmentSetup ) { this.components = [ new WorkspaceFolderComponent(this.workspaceFolderManager), @@ -69,7 +71,8 @@ export class ConfigurationDataProvider new EnvironmentComponent( this.featureManager, this.connectionManager, - this.configModel + this.configModel, + this.vpexEnvironmentSetup ), ]; this.disposables.push( diff --git a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts index 09f5ab9a1..7db4c6d34 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts @@ -4,6 +4,7 @@ import {BaseComponent} from "./BaseComponent"; import {ConfigurationTreeItem} from "./types"; import {ConnectionManager} from "../../configuration/ConnectionManager"; import {ConfigModel} from "../../configuration/models/ConfigModel"; +import {VpexEnvironmentSetup} from "../../language/VpexEnvironmentSetup"; const ENVIRONMENT_COMPONENT_ID = "ENVIRONMENT"; const getItemContext = (key: string, available: boolean) => @@ -13,12 +14,16 @@ export class EnvironmentComponent extends BaseComponent { constructor( private readonly featureManager: FeatureManager, private readonly connectionManager: ConnectionManager, - private readonly configModel: ConfigModel + private readonly configModel: ConfigModel, + private readonly vpexEnvironmentSetup: VpexEnvironmentSetup ) { super(); this.featureManager.onDidChangeState("environment.dependencies", () => this.onDidChangeEmitter.fire() ); + this.vpexEnvironmentSetup.onDidChangeState(() => + this.onDidChangeEmitter.fire() + ); } public async getRoot(): Promise { @@ -106,6 +111,27 @@ export class EnvironmentComponent extends BaseComponent { } } } + + // VPEX one-click setup entry: runs `databricks dbconnect init/sync` + // and auto-adopts the resulting .venv. + const vpexReady = this.vpexEnvironmentSetup.ready; + children.push({ + contextValue: getItemContext("vpex", vpexReady), + label: vpexReady + ? "Databricks Connect (VPEX): serverless-v4 · py3.12" + : "Set up Databricks Connect (VPEX)", + tooltip: vpexReady + ? "Environment ready: serverless-v4 (Python 3.12, databricks-connect 17.3). Click to re-run setup." + : "One-click setup of the Databricks Connect Python environment.", + iconPath: vpexReady + ? new ThemeIcon("check") + : new ThemeIcon("rocket"), + command: { + title: "Set up Python Environment (VPEX)", + command: "databricks.environment.setupVpex", + }, + }); + return children; } }