feat(python-setup): add the Python-setup preset picker (Full / DB Connect / Python)

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,

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<SetupLocalInvocation["compute"]> = [];
+        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");
+    });
+});

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<ResolvedCompute>;
 
+    /**
+     * 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<SetupPreset | undefined>;
+
     /**
      * 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

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<PresetPickItem>;
+    }) as <T extends QuickPickItem>() => QuickPick<T>;
+    return {create, created};
+}
+
 function makeWiring(
     overrides: Partial<PythonSetupWiringDeps> = {}
 ): 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);
+    });
+});