diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index b626341c5..3b819341f 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -15,6 +15,7 @@ import {ConnectionManager} from "./configuration/ConnectionManager"; import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider"; import {ClusterModel} from "./cluster/ClusterModel"; import {ClusterCommands} from "./cluster/ClusterCommands"; +import {Cluster} from "./sdk-extensions/Cluster"; import {ConfigurationDataProvider} from "./ui/configuration-view/ConfigurationDataProvider"; import {composePythonSetupEntry} from "./ui/configuration-view/pythonSetupEntry"; import {routeEnvironmentSetup} from "./language/pythonSetupRouting"; @@ -1024,6 +1025,32 @@ export async function activate( "databricks.connection.attachClusterQuickPick" ) ), + createQuickPick: (...args) => window.createQuickPick(...args), + // The preset picker's title names the resolved runtime. Reuse the + // attached cluster's already-loaded DBR when its id matches; + // otherwise (e.g. a cluster just picked inline, whose attach has not + // propagated yet) fetch it by id so the title always carries the + // runtime. Any failure degrades to a generic title, so it must not + // reject into the setup flow. + clusterDbrVersion: async (clusterId) => { + try { + const attached = connectionManager.cluster; + if (attached?.id === clusterId) { + return attached.dbrVersion; + } + const apiClient = connectionManager.apiClient; + if (apiClient === undefined) { + return undefined; + } + const cluster = await Cluster.fromClusterId( + apiClient, + clusterId + ); + return cluster.dbrVersion; + } catch { + return undefined; + } + }, setActiveInterpreter: async (interpreterPath, root) => { await pythonExtensionWrapper.api.environments.updateActiveEnvironmentPath( interpreterPath, diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts index 31703e2e2..b2c9088b8 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts @@ -17,6 +17,7 @@ import { } from "../models/fixtures/setupLocalResults"; import {PythonSetupErrorAction} from "../utils/errorMessages"; import {SetupLocalInvocation} from "../utils/setupLocalArgs"; +import {SetupPreset} from "../utils/pythonSetupPresetPicker"; import { PythonSetupAttempt, PythonSetupOutcomeReport, @@ -136,6 +137,10 @@ function makeDeps( // Mirror the production wrapper: hand the task a log sink and a // (never-cancelled) progress token. withProgress: async (_title, task) => task(() => {}, makeToken()), + // Default: the user picks the recommended Full preset, so a run happens + // with no skip flags (matching the pre-picker behaviour these tests were + // written against). Tests that exercise other presets override this. + pickSetupPreset: async () => "full", // Telemetry defaults to a no-op sink; tests that assert on events pass // a recorder instead. recordSetupAttempt: () => () => {}, @@ -1067,6 +1072,7 @@ describe("PythonSetupEnvironmentSetup telemetry", () => { targetType: "serverless", serverlessVersion: "5", mode: "default", + setupPreset: "full", // hasPyprojectToml defaults to true, so this is not greenfield. isGreenfield: false, // First run for the project: not yet ready. @@ -1592,3 +1598,100 @@ describe("PythonSetupEnvironmentSetup telemetry", () => { ]); }); }); + +describe("PythonSetupEnvironmentSetup.setup preset selection", () => { + /** Assert on the invocation the CLI was actually handed. */ + function invocationFor(preset: SetupPreset) { + const cli = makeCli(); + const setup = new PythonSetupEnvironmentSetup( + makeDeps({cli, pickSetupPreset: async () => preset}) + ); + return setup.setup().then(() => cli.calls[0]); + } + + it("runs the full preset with no skip flags", async () => { + const invocation = await invocationFor("full"); + expect(invocation.skipConstraints).to.equal(undefined); + expect(invocation.skipDbconnect).to.equal(undefined); + }); + + it("runs the dbconnect preset with --no-constraints only", async () => { + const invocation = await invocationFor("dbconnect"); + expect(invocation.skipConstraints).to.equal(true); + expect(invocation.skipDbconnect).to.equal(undefined); + }); + + it("runs the python preset with both --no-constraints and --no-dbconnect", async () => { + const invocation = await invocationFor("python"); + expect(invocation.skipConstraints).to.equal(true); + expect(invocation.skipDbconnect).to.equal(true); + }); + + it("passes the resolved compute to the preset picker", async () => { + const seen: Array = []; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + resolveCompute: async () => ({ + status: "ok", + compute: {kind: "cluster", clusterId: "0710-abc"}, + }), + pickSetupPreset: async (compute) => { + seen.push(compute); + return "full"; + }, + }) + ); + + await setup.setup(); + + expect(seen).to.deep.equal([{kind: "cluster", clusterId: "0710-abc"}]); + }); + + it("does not run the CLI or record an attempt when the picker is dismissed", async () => { + const cli = makeCli(); + const telemetry = makeTelemetryRecorder(); + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + ...telemetry, + pickSetupPreset: async () => undefined, + }) + ); + + await setup.setup(); + + // Dismissing the picker is a deliberate bail-out before any run starts, + // so nothing is spawned and no attempt is recorded (mirroring a + // dismissed serverless-version prompt). + expect(cli.calls).to.have.length(0); + expect(telemetry.attempts).to.have.length(0); + expect(telemetry.results).to.have.length(0); + expect(setup.ready).to.equal(false); + }); + + it("records the chosen preset on the attempt (dbconnect keeps mode default)", async () => { + const telemetry = makeTelemetryRecorder(); + const setup = new PythonSetupEnvironmentSetup( + makeDeps({...telemetry, pickSetupPreset: async () => "dbconnect"}) + ); + + await setup.setup(); + + expect(telemetry.attempts[0].setupPreset).to.equal("dbconnect"); + // dbconnect skips only the pins (databricks-connect stays), so the + // legacy mode dimension is still "default". + expect(telemetry.attempts[0].mode).to.equal("default"); + }); + + it("maps the python preset to the constraints-only telemetry mode", async () => { + const telemetry = makeTelemetryRecorder(); + const setup = new PythonSetupEnvironmentSetup( + makeDeps({...telemetry, pickSetupPreset: async () => "python"}) + ); + + await setup.setup(); + + expect(telemetry.attempts[0].setupPreset).to.equal("python"); + expect(telemetry.attempts[0].mode).to.equal("constraints-only"); + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts index 99391faaa..2ff4f21b9 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts @@ -29,6 +29,7 @@ import { } from "../utils/reportSetupIssue"; import {isReauthRequiredError} from "../utils/authErrors"; import {SetupLocalInvocation} from "../utils/setupLocalArgs"; +import {presetToFlags, SetupPreset} from "../utils/pythonSetupPresetPicker"; import { PythonSetupAttempt, PythonSetupResultReporter, @@ -123,6 +124,18 @@ export interface PythonSetupSetupDeps { */ resolveCompute: () => Promise; + /** + * Ask the user which preset to provision for the resolved compute (Full / + * DB Connect / Python). Returns the chosen preset, or `undefined` when the + * picker is dismissed — a deliberate bail-out before any run starts, so the + * flow stops silently and records no attempt, mirroring a dismissed + * serverless-version prompt. Called after compute resolves so the picker's + * title can name the resolved target. + */ + pickSetupPreset: ( + compute: SetupCompute + ) => Promise; + /** * Point the MS Python extension at the provisioned venv interpreter for * `projectRoot`. The root is passed in (not re-read) so adoption always @@ -368,8 +381,17 @@ export class PythonSetupEnvironmentSetup implements Disposable { } const compute = resolved.compute; + // A dismissed picker is a deliberate bail-out before any run starts, so + // return silently and record no attempt — like the dismissed + // serverless-version prompt above. + const preset = await this.deps.pickSetupPreset(compute); + if (preset === undefined) { + return; + } + const invocation: SetupLocalInvocation = { compute, + ...presetToFlags(preset), }; // From here a run really happens, so the attempt is recorded and every @@ -378,7 +400,8 @@ export class PythonSetupEnvironmentSetup implements Disposable { // spawn and interpreter adoption. const {reportResult, packageManager} = await this.recordAttempt( invocation, - cwd + cwd, + preset ); // Per-run report context: the static build info plus this run's manager. const reportEnv: ReportEnvironment = { @@ -593,7 +616,8 @@ export class PythonSetupEnvironmentSetup implements Disposable { */ private async recordAttempt( invocation: SetupLocalInvocation, - projectRoot: string + projectRoot: string, + setupPreset: SetupPreset ): Promise<{ reportResult: PythonSetupResultReporter; packageManager: PrimaryManager; @@ -635,8 +659,11 @@ export class PythonSetupEnvironmentSetup implements Disposable { serverlessVersion: compute.kind === "serverless" ? compute.version : undefined, // --no-dbconnect is the orthogonal spelling of the legacy - // --constraints-only, so it maps to that telemetry mode. + // --constraints-only, so it maps to that telemetry mode. The + // richer, unambiguous axis is `setupPreset`; `mode` is kept for + // dashboard continuity. mode: invocation.skipDbconnect ? "constraints-only" : "default", + setupPreset, isGreenfield, // A run against a project already marked ready this session is a // re-run (the ready row's Re-run button / row click); anything diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts index 530ae33f0..2be23a135 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -1,11 +1,12 @@ import {expect} from "chai"; -import {commands, env, Uri, window} from "vscode"; +import {commands, env, QuickPick, QuickPickItem, Uri, window} from "vscode"; import { makePythonSetupDeps, makePythonSetupVisibility, PythonSetupWiringDeps, resolveComputeFrom, } from "./pythonSetupDeps"; +import {PresetPickItem} from "../utils/pythonSetupPresetPicker"; import {PythonSetupState} from "../../vscode-objs/StateStorage"; import {SetupCompute} from "./PythonSetupEnvironmentSetup"; import {Telemetry} from "../../telemetry"; @@ -165,6 +166,64 @@ describe("resolveComputeFrom", () => { }); }); +/** + * A minimal scriptable QuickPick stand-in (see `AiToolsCommands.test.ts`). + * `onShow` decides which item is selected, or dismisses. + */ +class FakeQuickPick { + title?: string; + placeholder?: string; + items: readonly QuickPickItem[] = []; + selectedItems: readonly QuickPickItem[] = []; + private acceptCbs: Array<() => void> = []; + private hideCbs: Array<() => void> = []; + constructor( + private readonly onShow: ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" + ) {} + onDidAccept(cb: () => void) { + this.acceptCbs.push(cb); + return {dispose() {}}; + } + onDidHide(cb: () => void) { + this.hideCbs.push(cb); + return {dispose() {}}; + } + show() { + const r = this.onShow(this); + if (r === "dismiss") { + this.hideCbs.forEach((cb) => cb()); + return; + } + this.selectedItems = r.selected; + this.acceptCbs.forEach((cb) => cb()); + } + hide() { + this.hideCbs.forEach((cb) => cb()); + } + dispose() {} +} + +/** + * A `createQuickPick` factory (typed as the real `window.createQuickPick` + * seam) that scripts the widget's outcome and records the created instances so + * a test can read the title the wiring set. + */ +function fakeCreateQuickPick( + onShow: ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" +) { + const created: FakeQuickPick[] = []; + const create = (() => { + const pick = new FakeQuickPick(onShow); + created.push(pick); + return pick as unknown as QuickPick; + }) as () => QuickPick; + return {create, created}; +} + function makeWiring( overrides: Partial = {} ): PythonSetupWiringDeps { @@ -185,6 +244,8 @@ function makeWiring( promptServerlessVersion: async () => "4", persistServerlessVersion: async () => {}, promptSelectCompute: async () => undefined, + createQuickPick: fakeCreateQuickPick(() => "dismiss").create, + clusterDbrVersion: async () => undefined, setActiveInterpreter: async () => {}, persistSetupState: () => {}, log: {append: () => {}, show: () => {}}, @@ -1041,3 +1102,58 @@ describe("makePythonSetupDeps showReauthPrompt", () => { expect(executed).to.have.length(0); }); }); + +describe("makePythonSetupDeps pickSetupPreset", () => { + it("titles the picker with the serverless target and returns the picked preset", async () => { + const {create, created} = fakeCreateQuickPick((pick) => ({ + // Pick the DB Connect row (the second one). + selected: [pick.items[1]], + })); + const deps = makePythonSetupDeps(makeWiring({createQuickPick: create})); + + const preset = await deps.pickSetupPreset({ + kind: "serverless", + version: "5", + }); + + expect(preset).to.equal("dbconnect"); + expect(created[0].title).to.equal( + "Set up Python environment for serverless v5" + ); + }); + + it("titles the picker with the cluster's runtime, resolved from its DBR", async () => { + const requested: string[] = []; + const {create, created} = fakeCreateQuickPick(() => "dismiss"); + const deps = makePythonSetupDeps( + makeWiring({ + createQuickPick: create, + clusterDbrVersion: async (id) => { + requested.push(id); + return [17, 3, "x"]; + }, + }) + ); + + await deps.pickSetupPreset({kind: "cluster", clusterId: "0710-abc"}); + + // The DBR is looked up for the resolved cluster id, and its major.minor + // becomes the runtime shown in the title. + expect(requested).to.deep.equal(["0710-abc"]); + expect(created[0].title).to.equal( + "Set up Python environment for Runtime 17.3" + ); + }); + + it("returns undefined when the picker is dismissed", async () => { + const deps = makePythonSetupDeps( + makeWiring({ + createQuickPick: fakeCreateQuickPick(() => "dismiss").create, + }) + ); + + expect( + await deps.pickSetupPreset({kind: "serverless", version: "5"}) + ).to.equal(undefined); + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index e9300fd09..adbf317f8 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -11,6 +11,10 @@ import {PythonSetupErrorAction} from "../utils/errorMessages"; import {ReportEnvironment} from "../utils/reportSetupIssue"; import {isUvSetupSuitable} from "../utils/pythonSetupGate"; import {withElapsedProgress} from "../utils/setupProgress"; +import { + computeTargetLabel, + pickSetupPreset, +} from "../utils/pythonSetupPresetPicker"; import {formatSetupLog, formatSetupNotification} from "../utils/setupSummary"; import {venvInterpreterPath} from "../utils/venvInterpreterPath"; import {readVenvProjectName} from "../utils/venvProjectName"; @@ -159,6 +163,22 @@ export interface PythonSetupWiringDeps { * propagates asynchronously, so an immediate re-read would race it. */ promptSelectCompute: () => Promise; + /** + * The explicit-lifecycle QuickPick factory (`window.createQuickPick`), + * injected so the preset picker can be driven in tests without a VS Code + * host. + */ + createQuickPick: (typeof window)["createQuickPick"]; + /** + * The DBR version of the attached cluster, as `[major, minor, patch]` (see + * `Cluster.dbrVersion`), for the preset picker's title. `undefined` when it + * cannot be resolved (unknown cluster, custom image), in which case the + * title falls back to a generic cluster label. Only consulted for a cluster + * target. + */ + clusterDbrVersion: ( + clusterId: string + ) => Promise | undefined>; /** Point the MS Python extension at an interpreter path (project-scoped). */ setActiveInterpreter: (interpreterPath: string, root: Uri) => Promise; /** Persist the post-setup state (workspace-scoped) for drift detection. */ @@ -238,6 +258,18 @@ export function makePythonSetupDeps( } return {status: "ok", compute: {kind: "serverless", version}}; }, + pickSetupPreset: async (compute) => { + // A cluster's runtime label needs its DBR version; serverless + // carries its version directly, so no lookup is needed there. + const dbrVersion = + compute.kind === "cluster" + ? await wiring.clusterDbrVersion(compute.clusterId) + : undefined; + return pickSetupPreset( + computeTargetLabel(compute, dbrVersion), + wiring.createQuickPick + ); + }, adoptInterpreter: async (venvPath: string, projectRoot: string) => { await wiring.setActiveInterpreter( venvInterpreterPath(venvPath), diff --git a/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.test.ts b/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.test.ts new file mode 100644 index 000000000..7bc609dec --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.test.ts @@ -0,0 +1,215 @@ +import {expect} from "chai"; +import {QuickPick, QuickPickItem} from "vscode"; +import { + buildPresetPickItems, + computeTargetLabel, + pickSetupPreset, + presetToFlags, + PresetPickItem, +} from "./pythonSetupPresetPicker"; + +/** + * A minimal, scriptable stand-in for a VS Code QuickPick, mirroring the fake + * used in `AiToolsCommands.test.ts`. `onAccept` decides which item is selected + * (from the items assigned to the pick) and whether the pick is accepted or + * dismissed, then drives the accept/hide callbacks the way the real widget + * would. + */ +class FakeQuickPick { + title?: string; + placeholder?: string; + items: readonly QuickPickItem[] = []; + selectedItems: readonly QuickPickItem[] = []; + private acceptCbs: Array<() => void> = []; + private hideCbs: Array<() => void> = []; + public disposed = false; + + constructor( + private readonly onAccept: ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" + ) {} + + onDidAccept(cb: () => void) { + this.acceptCbs.push(cb); + return {dispose() {}}; + } + onDidHide(cb: () => void) { + this.hideCbs.push(cb); + return {dispose() {}}; + } + show() { + const result = this.onAccept(this); + if (result === "dismiss") { + this.hideCbs.forEach((cb) => cb()); + return; + } + this.selectedItems = result.selected; + this.acceptCbs.forEach((cb) => cb()); + } + hide() { + this.hideCbs.forEach((cb) => cb()); + } + dispose() { + this.disposed = true; + } +} + +/** + * A factory that hands `pickSetupPreset` a `FakeQuickPick`, typed as the real + * `window.createQuickPick` seam. `onAccept` scripts the widget's outcome, and + * `created` exposes the instance so a test can read the title/placeholder/items + * the picker set on it. + */ +function fakeCreateQuickPick( + onAccept: ( + pick: FakeQuickPick + ) => {selected: readonly QuickPickItem[]} | "dismiss" +) { + const created: FakeQuickPick[] = []; + const create = (() => { + const pick = new FakeQuickPick(onAccept); + created.push(pick); + return pick as unknown as QuickPick; + }) as () => QuickPick; + return {create, created}; +} + +describe("presetToFlags", () => { + it("maps the full preset to no skip flags", () => { + expect(presetToFlags("full")).to.deep.equal({}); + }); + + it("maps the dbconnect preset to --no-constraints only", () => { + expect(presetToFlags("dbconnect")).to.deep.equal({ + skipConstraints: true, + }); + }); + + it("maps the python preset to both --no-constraints and --no-dbconnect", () => { + expect(presetToFlags("python")).to.deep.equal({ + skipConstraints: true, + skipDbconnect: true, + }); + }); +}); + +describe("buildPresetPickItems", () => { + it("lists the three presets in the Full, DB Connect, Python order", () => { + const items = buildPresetPickItems(); + expect(items.map((i) => i.preset)).to.deep.equal([ + "full", + "dbconnect", + "python", + ]); + }); + + it("presents the Full row with its verbatim copy, starred and first", () => { + const [full] = buildPresetPickItems(); + expect(full.label).to.equal("$(star-full) Full environment setup"); + expect(full.description).to.equal( + "Recommended so code runs as it would on Databricks" + ); + expect(full.detail).to.equal( + "Installs matching Python + Databricks Connect versions and pins cluster dependencies." + ); + }); + + it("presents the DB Connect row with its verbatim copy", () => { + const dbconnect = buildPresetPickItems()[1]; + expect(dbconnect.label).to.equal("$(tools) DB Connect setup"); + expect(dbconnect.description).to.equal( + "Recommended to run Spark code remotely." + ); + expect(dbconnect.detail).to.equal( + "Installs matching Python + Databricks Connect versions only." + ); + }); + + it("presents the Python row with its verbatim copy and no description", () => { + const python = buildPresetPickItems()[2]; + expect(python.label).to.equal("$(code) Python setup"); + expect(python.description).to.equal(undefined); + expect(python.detail).to.equal( + "Installs matching Python version only." + ); + }); +}); + +describe("computeTargetLabel", () => { + it("labels a serverless target with its vN version", () => { + expect(computeTargetLabel({kind: "serverless", version: "5"})).to.equal( + "serverless v5" + ); + }); + + it("labels a cluster target with its major.minor runtime", () => { + expect( + computeTargetLabel({kind: "cluster", clusterId: "0710-abc"}, [ + 17, + 3, + "x", + ]) + ).to.equal("Runtime 17.3"); + }); + + it("falls back to a generic cluster label when the runtime is unparsable", () => { + expect( + computeTargetLabel({kind: "cluster", clusterId: "0710-abc"}, [ + "x", + "x", + "x", + ]) + ).to.equal("the attached cluster"); + }); + + it("falls back to a generic cluster label when no runtime is known", () => { + expect( + computeTargetLabel({kind: "cluster", clusterId: "0710-abc"}) + ).to.equal("the attached cluster"); + }); +}); + +describe("pickSetupPreset", () => { + it("titles the picker with the resolved compute and documents the side effects", async () => { + const {create, created} = fakeCreateQuickPick(() => "dismiss"); + + await pickSetupPreset("serverless v5", create); + + expect(created[0].title).to.equal( + "Set up Python environment for serverless v5" + ); + expect(created[0].placeholder).to.equal( + "Create a uv managed .venv and pyproject.toml (if one exists, it's saved to pyproject.toml.bak)" + ); + // The picker offers exactly the three preset rows. + expect(created[0].items).to.have.length(3); + }); + + it("resolves the chosen preset when the user accepts a row", async () => { + const {create} = fakeCreateQuickPick((pick) => ({ + // Accept the DB Connect row (the second one). + selected: [pick.items[1]], + })); + + expect(await pickSetupPreset("Runtime 17.3", create)).to.equal( + "dbconnect" + ); + }); + + it("resolves undefined when the user dismisses the picker", async () => { + const {create} = fakeCreateQuickPick(() => "dismiss"); + + expect(await pickSetupPreset("Runtime 17.3", create)).to.equal( + undefined + ); + }); + + it("disposes the picker once it hides", async () => { + const {create, created} = fakeCreateQuickPick(() => "dismiss"); + + await pickSetupPreset("serverless v5", create); + + expect(created[0].disposed).to.equal(true); + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.ts b/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.ts new file mode 100644 index 000000000..dad6b40b9 --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/pythonSetupPresetPicker.ts @@ -0,0 +1,143 @@ +import {QuickPick, QuickPickItem} from "vscode"; +import {SetupLocalInvocation} from "./setupLocalArgs"; + +/** The resolved compute target a setup run provisions for. */ +export type SetupComputeTarget = SetupLocalInvocation["compute"]; + +/** + * The Python-setup preset the user picks. Each tier nests inside the previous + * one, so they map onto the two orthogonal `setup-local` skip flags (see + * {@link presetToFlags}): + * + * - `full` — matching Python + Databricks Connect, and the runtime/dependency + * pins. No flags. + * - `dbconnect` — matching Python + Databricks Connect, but no pins + * (`--no-constraints`). + * - `python` — matching Python only (`--no-constraints --no-dbconnect`). + * + * The fourth flag combination (pins without databricks-connect, today's + * `--constraints-only`) is intentionally not a tier: keeping the tiers strictly + * nested is what lets a single-select picker present them as one axis. + */ +export type SetupPreset = "full" | "dbconnect" | "python"; + +/** The `setup-local` skip flags a preset resolves to. */ +export interface SetupPresetFlags { + skipConstraints?: boolean; + skipDbconnect?: boolean; +} + +/** + * Map a preset to the orthogonal skip flags the CLI invocation carries. Only + * the `true` flags are set, so a `full` run adds nothing and the argv stays + * flag-free (matching `buildSetupLocalArgs`, which pushes a flag only when its + * field is truthy). + */ +export function presetToFlags(preset: SetupPreset): SetupPresetFlags { + switch (preset) { + case "full": + return {}; + case "dbconnect": + return {skipConstraints: true}; + case "python": + return {skipConstraints: true, skipDbconnect: true}; + } +} + +/** + * The human label for the resolved compute target, used in the picker title + * (e.g. "Set up Python environment for Runtime 17.3"). Serverless carries its + * version directly; a cluster is described by its DBR major.minor when the + * `dbrVersion` is known and parseable, and otherwise by a generic label — the + * cluster id is never a resolved runtime and a cluster *name* is user-chosen + * (routinely a person's name), so neither is shown. + */ +export function computeTargetLabel( + compute: SetupComputeTarget, + dbrVersion?: Array +): string { + if (compute.kind === "serverless") { + return `serverless v${compute.version}`; + } + const [major, minor] = dbrVersion ?? []; + if (typeof major === "number" && typeof minor === "number") { + return `Runtime ${major}.${minor}`; + } + return "the attached cluster"; +} + +/** A preset QuickPick row; `preset` is the choice the row resolves to. */ +export interface PresetPickItem extends QuickPickItem { + preset: SetupPreset; +} + +/** Documents the shared side effects every tier has (all provision a venv). */ +const PLACEHOLDER = + "Create a uv managed .venv and pyproject.toml (if one exists, it's saved to pyproject.toml.bak)"; + +/** + * Build the three preset rows in nesting order (Full → DB Connect → Python). + * Pure, so the copy is unit-testable without a VS Code host. Full is listed + * first and starred as the recommendation; the codicons (`$(...)`) render as + * inline icons in the QuickPick. + */ +export function buildPresetPickItems(): PresetPickItem[] { + return [ + { + label: "$(star-full) Full environment setup", + description: "Recommended so code runs as it would on Databricks", + detail: "Installs matching Python + Databricks Connect versions and pins cluster dependencies.", + preset: "full", + }, + { + label: "$(tools) DB Connect setup", + description: "Recommended to run Spark code remotely.", + detail: "Installs matching Python + Databricks Connect versions only.", + preset: "dbconnect", + }, + { + label: "$(code) Python setup", + detail: "Installs matching Python version only.", + preset: "python", + }, + ]; +} + +/** + * Show the single-select preset picker for a resolved compute target and return + * the chosen preset, or `undefined` if the user dismissed it. + * + * Uses the explicit-lifecycle `createQuickPick` API (not the one-shot + * `showQuickPick`) so `onDidAccept` resolves the chosen row and `onDidHide` + * reports a dismissal, matching the ticket's UX contract. `createQuickPick` is + * injected (the wiring passes `window.createQuickPick`) so the flow is + * unit-testable without a VS Code host. The title carries `computeLabel` (the + * resolved compute, e.g. "Runtime 17.3" or "serverless v5") and the placeholder + * documents the shared side effects. + */ +export function pickSetupPreset( + computeLabel: string, + createQuickPick: () => QuickPick +): Promise { + const quickPick = createQuickPick(); + quickPick.title = `Set up Python environment for ${computeLabel}`; + quickPick.placeholder = PLACEHOLDER; + quickPick.items = buildPresetPickItems(); + + return new Promise((resolve) => { + let picked: SetupPreset | undefined; + quickPick.onDidAccept(() => { + const selected = quickPick.selectedItems[0]; + if (selected === undefined) { + return; + } + picked = selected.preset; + quickPick.hide(); + }); + quickPick.onDidHide(() => { + resolve(picked); + quickPick.dispose(); + }); + quickPick.show(); + }); +} diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index 1e9b2a322..8e0c91736 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -154,6 +154,11 @@ import type { PythonSetupErrorCode, } from "../python-setup/models/PythonSetupResult"; export type {PythonSetupMode, PythonSetupErrorCode}; +// The setup preset is an extension-side concept (the picker's tiers), not part +// of the CLI wire contract, so it is owned by the picker util. Type-only, for +// the attempt event's schema. +import type {SetupPreset} from "../python-setup/utils/pythonSetupPresetPicker"; +export type {SetupPreset}; /** * How a setup run ended. @@ -541,6 +546,7 @@ export class EventTypes { targetType: ComputeType; serverlessVersion?: string; mode: PythonSetupMode; + setupPreset: SetupPreset; isGreenfield?: boolean; trigger: PythonSetupRunTrigger; }> = { @@ -570,6 +576,14 @@ export class EventTypes { comment: "Whether databricks-connect is included (default) or only the runtime constraints (constraints-only)", }, + setupPreset: { + comment: + "The preset tier the user picked: full (matching Python + Databricks Connect + " + + "pinned cluster dependencies), dbconnect (matching Python + Databricks Connect, " + + "no pins), or python (matching Python only). Disambiguates the two orthogonal " + + "skip axes that the two-value mode field cannot: a dbconnect run skips the pins " + + "yet reports mode=default", + }, isGreenfield: { comment: "Whether the project has no pyproject.toml yet. Omitted unless the project is " + diff --git a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts index 24696dfda..641e934f0 100644 --- a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts +++ b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts @@ -40,6 +40,7 @@ describe(__filename, () => { targetType: "serverless", serverlessVersion: "5", mode: "default", + setupPreset: "full", isGreenfield: true, trigger: "initial", }); @@ -58,6 +59,7 @@ describe(__filename, () => { "event.targetType": "serverless", "event.serverlessVersion": "5", "event.mode": "default", + "event.setupPreset": "full", "event.isGreenfield": "true", "event.trigger": "initial", }); @@ -75,6 +77,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", }); reportResult({outcome: "ok"}); @@ -96,6 +99,7 @@ describe(__filename, () => { packageManager: "pip", targetType: "cluster", mode: "constraints-only", + setupPreset: "python", serverlessVersion: undefined, isGreenfield: undefined, trigger: "initial", @@ -114,6 +118,7 @@ describe(__filename, () => { "event.packageManager": "pip", "event.targetType": "cluster", "event.mode": "constraints-only", + "event.setupPreset": "python", "event.trigger": "initial", }); expect(events[1].props).to.deep.equal({ @@ -134,6 +139,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", isGreenfield: false, trigger: "initial", }); @@ -162,6 +168,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", }); reportResult({ @@ -181,6 +188,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", }); reportResult({outcome: "failed", failurePhase: "provision"}); @@ -195,6 +203,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", }); reportResult({outcome: "ok"}); @@ -221,6 +230,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", })({outcome: "ok", pythonSetupFlow}); @@ -244,6 +254,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", })({outcome: "ok", envKey}); expect(events[1].props["event.envKey"]).to.equal(envKey); @@ -272,6 +283,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", })({outcome: "ok", envKey}); expect(events[1].props["event.envKey"]).to.equal("other"); @@ -286,6 +298,7 @@ describe(__filename, () => { targetType: "serverless", serverlessVersion: "5", mode: "default", + setupPreset: "full", trigger: "initial", })({ outcome: "ok", @@ -327,6 +340,7 @@ describe(__filename, () => { targetType: "serverless", serverlessVersion: "5", mode: "default", + setupPreset: "full", trigger: "initial", })({outcome: "ok", envKey: "serverless/serverless-v5", warnings: []}); @@ -344,6 +358,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", })({ outcome: "ok", @@ -380,6 +395,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", })({outcome: "cancelled"}); @@ -399,6 +415,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", clusterId: "0710-142042-secretcluster", projectPath: "/Users/jane/projects/acme", @@ -417,6 +434,7 @@ describe(__filename, () => { expect(Object.keys(events[0].props).sort()).to.deep.equal([ "event.mode", "event.packageManager", + "event.setupPreset", "event.targetType", "event.trigger", "version", @@ -454,6 +472,7 @@ describe(__filename, () => { packageManager: "uv", targetType: "cluster", mode: "default", + setupPreset: "full", trigger: "initial", }); diff --git a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts index 68c03066f..a7e9a89dd 100644 --- a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts +++ b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts @@ -11,6 +11,7 @@ import type { PythonSetupOptOutSource, PythonSetupOutcome, PythonSetupRunTrigger, + SetupPreset, TargetCompute, } from "./constants"; import {PythonSetupWarning} from "../python-setup/models/PythonSetupResult"; @@ -26,6 +27,13 @@ export interface PythonSetupAttempt { /** The chosen serverless environment version; absent for clusters. */ serverlessVersion?: string; mode: PythonSetupMode; + /** + * The preset tier the user picked (full | dbconnect | python). Kept + * alongside {@link mode} because the two-value mode field cannot represent + * the orthogonal skip axes the picker enables: a `dbconnect` run skips the + * pins yet still reports `mode: "default"`. + */ + setupPreset: SetupPreset; /** * Whether the project has no `pyproject.toml` yet, or `undefined` when the * signal would be misleading — for a pip/conda project the absence of a @@ -276,6 +284,7 @@ Telemetry.prototype.recordPythonSetupAttempt = function ( packageManager: attempt.packageManager, targetType: attempt.targetType, mode: attempt.mode, + setupPreset: attempt.setupPreset, trigger: attempt.trigger, ...(attempt.serverlessVersion !== undefined ? {serverlessVersion: attempt.serverlessVersion} diff --git a/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts b/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts index e059e73b4..c1edda213 100644 --- a/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts @@ -206,11 +206,18 @@ describe("Set up local Python environment (uv) on serverless", async function () it("should set up the environment with uv (setup-local)", async () => { // The router command; a clean (uv-suitable) project routes it to the uv // flow, which shells out to `databricks environments setup-local`. - // Compute + version are already resolved, so no prompt. + // Compute + version are already resolved, so the preset picker below is + // the only prompt. await executeCommandWhenAvailable( "Databricks: Setup python environment" ); + // Once compute is resolved the flow asks which preset to provision; take + // the recommended Full tier (starred, first) so the full environment is + // set up (matching Python + Databricks Connect + pinned dependencies). + const presetInput = await waitForQuickInput(); + await presetInput.selectQuickPick(0); + await waitForVenvInterpreter(projectDir); // The completion toast is informational and can expire before a poll @@ -262,12 +269,22 @@ describe("Set up local Python environment (uv) on serverless", async function () // executeCommand resolves when the (re-entrancy-guarded) run settles; a // re-run over an already-provisioned env is a warm uv sync, so it returns // quickly rather than paying another cold provision. - await browser.executeWorkbench(async (vscode) => { - await vscode.commands.executeCommand( + // rerunPythonEnv is palette-hidden (when:false), so fire it by id — but + // do NOT await it inside executeWorkbench: the command now blocks on the + // preset picker (driven just below), so awaiting here would deadlock — + // the callback can't return while the picker is open, so the picker + // could never be answered. + await browser.executeWorkbench((vscode) => { + void vscode.commands.executeCommand( "databricks.environment.rerunPythonEnv" ); }); + // Take the recommended Full tier again so the re-run provisions the same + // environment. + const rerunPresetInput = await waitForQuickInput(); + await rerunPresetInput.selectQuickPick(0); + // A fresh "Python environment ready" toast is the re-run-SPECIFIC success // signal. A failed run never clears readiness (readyRoots/state persist) // and never removes the .venv, so the interpreter and the ready row alone