diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index 626557093..0ca8d54d3 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -3,11 +3,57 @@ import {ApiClient} from "@databricks/sdk-experimental"; import {Cluster} from "../sdk-extensions"; import assert from "assert"; import {mock} from "ts-mockito"; +import {QuickPickItem, window} from "vscode"; import { + ConnectionCommands, formatClusterState, formatQuickPickClusterDetails, } from "./ConnectionCommands"; +/** + * Scriptable stand-in for the compute QuickPick. The test drives + * {@link accept}/{@link hide} explicitly, since the command registers its + * handlers only after `show()`. + */ +class FakeComputeQuickPick { + title?: string; + keepScrollPosition = false; + busy = false; + canSelectMany = false; + items: readonly QuickPickItem[] = []; + selectedItems: readonly QuickPickItem[] = []; + disposed = false; + private readonly acceptCbs: Array<() => void | Promise> = []; + private readonly hideCbs: Array<() => void> = []; + + onDidAccept(cb: () => void | Promise) { + this.acceptCbs.push(cb); + return {dispose() {}}; + } + onDidHide(cb: () => void) { + this.hideCbs.push(cb); + return {dispose() {}}; + } + show() {} + dispose() { + this.disposed = true; + } + + /** Simulate the user picking `selected` and pressing Enter. */ + async accept(selected: readonly QuickPickItem[]) { + this.selectedItems = selected; + for (const cb of this.acceptCbs) { + await cb(); + } + } + /** Simulate the user dismissing the picker (Escape). */ + async hide() { + for (const cb of this.hideCbs) { + cb(); + } + } +} + describe(__filename, () => { it("attach cluster quickpick: correctly format cluster details", () => { const clusterDetails = formatQuickPickClusterDetails( @@ -34,4 +80,157 @@ describe(__filename, () => { assert.equal(formatClusterState("ERROR"), "Error"); assert.equal(formatClusterState("UNKNOWN"), "Unknown"); }); + + describe("attachClusterQuickPick return contract", () => { + let originalCreateQuickPick: typeof window.createQuickPick; + let fakePick: FakeComputeQuickPick; + let attachCalls: string[]; + let commands: ConnectionCommands; + + beforeEach(() => { + originalCreateQuickPick = window.createQuickPick; + fakePick = new FakeComputeQuickPick(); + (window as unknown as {createQuickPick: unknown}).createQuickPick = + () => fakePick; + + attachCalls = []; + const connectionManager = { + workspaceClient: {}, + databricksWorkspace: {userName: "me"}, + attachCluster: async (id: string) => { + attachCalls.push(id); + }, + enableServerless: async () => {}, + }; + const clusterModel = { + refresh() {}, + onDidChange() { + return {dispose() {}}; + }, + roots: [], + }; + commands = new ConnectionCommands( + {} as never, + connectionManager as never, + clusterModel as never, + {} as never, + {} as never, + {} as never + ); + }); + + afterEach(() => { + (window as unknown as {createQuickPick: unknown}).createQuickPick = + originalCreateQuickPick; + }); + + const clusterItem = (id: string) => + ({ + label: `cluster ${id}`, + cluster: {id}, + }) as unknown as QuickPickItem; + + it("resolves to the attached cluster and attaches it exactly once", async () => { + const resultP = commands.attachClusterQuickPickCommand()(); + await fakePick.accept([clusterItem("c1")]); + + assert.deepEqual(await resultP, { + kind: "cluster", + clusterId: "c1", + }); + assert.deepEqual(attachCalls, ["c1"]); + }); + + it("resolves to undefined when the picker is dismissed", async () => { + const resultP = commands.attachClusterQuickPickCommand()(); + await fakePick.hide(); + + assert.equal(await resultP, undefined); + assert.deepEqual(attachCalls, []); + }); + + it("guards re-entry: a second Enter neither re-attaches nor changes the result", async () => { + // Without the `settled` guard a double Enter would attach twice. + const resultP = commands.attachClusterQuickPickCommand()(); + await fakePick.accept([clusterItem("c1")]); + await fakePick.accept([clusterItem("c2")]); + + assert.deepEqual(await resultP, { + kind: "cluster", + clusterId: "c1", + }); + assert.deepEqual(attachCalls, ["c1"]); + }); + + it("resolves to undefined when accepted with nothing highlighted", async () => { + const resultP = commands.attachClusterQuickPickCommand()(); + await fakePick.accept([]); + + assert.equal(await resultP, undefined); + assert.deepEqual(attachCalls, []); + }); + + it("still returns the chosen cluster when the attach silently fails (best-effort)", async () => { + // attachCluster is @onError(throw:false): a failed write resolves + // without throwing, so we still return the picked target by design. + const attach = async () => {}; + const cmds = new ConnectionCommands( + {} as never, + { + workspaceClient: {}, + databricksWorkspace: {userName: "me"}, + attachCluster: attach, + enableServerless: async () => {}, + } as never, + { + refresh() {}, + onDidChange() { + return {dispose() {}}; + }, + roots: [], + } as never, + {} as never, + {} as never, + {} as never + ); + + const resultP = cmds.attachClusterQuickPickCommand()(); + await fakePick.accept([clusterItem("c1")]); + + assert.deepEqual(await resultP, { + kind: "cluster", + clusterId: "c1", + }); + }); + + it("resolves to undefined if an unexpected error is thrown during selection", async () => { + // A throw from any step must settle the Promise, not hang. + const cmds = new ConnectionCommands( + {} as never, + { + workspaceClient: {}, + databricksWorkspace: {userName: "me"}, + attachCluster: async () => { + throw new Error("unexpected"); + }, + enableServerless: async () => {}, + } as never, + { + refresh() {}, + onDidChange() { + return {dispose() {}}; + }, + roots: [], + } as never, + {} as never, + {} as never, + {} as never + ); + + const resultP = cmds.attachClusterQuickPickCommand()(); + await fakePick.accept([clusterItem("c1")]); + + assert.equal(await resultP, undefined); + }); + }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index b930d2483..2ee19cfe3 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -1,5 +1,5 @@ import {Cluster} from "../sdk-extensions"; -import {compute} from "@databricks/sdk-experimental"; +import {compute, logging} from "@databricks/sdk-experimental"; import { Disposable, QuickPickItem, @@ -30,7 +30,19 @@ import { } from "../python-setup/utils/serverlessVersionResolver"; import {pickServerlessVersion} from "../python-setup/utils/serverlessVersionPicker"; import {collectServerlessVersionObservations} from "../python-setup/utils/serverlessVersionObservations"; +import type {SetupCompute} from "../python-setup/controllers/PythonSetupEnvironmentSetup"; import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager"; +import {Loggers} from "../logger"; + +// eslint-disable-next-line @typescript-eslint/naming-convention +const {NamedLogger} = logging; + +/** + * A compute target picked in the QuickPick. Aliased to the `setup-local` compute + * shape so a drift is a compile error, not a silent mismatch through + * `executeCommand`'s untyped generic. Serverless is always version-complete. + */ +export type SelectedCompute = SetupCompute; function formatQuickPickClusterSize(sizeInMB: number): string { if (sizeInMB > 1024) { @@ -143,8 +155,14 @@ export class ConnectionCommands implements Disposable { }; } + /** + * Resolves once the picker closes to the attached compute, or `undefined` on + * dismissal / "Create New Cluster". The return lets python-setup use the + * selection directly rather than re-reading the connection manager (whose + * cluster attach is async and would race). Other callers ignore it. + */ attachClusterQuickPickCommand() { - return async (title?: string) => { + return async (title?: string): Promise => { const workspaceClient = this.connectionManager.workspaceClient; const me = this.connectionManager.databricksWorkspace?.userName; if (!workspaceClient || !me) { @@ -204,33 +222,73 @@ export class ConnectionCommands implements Disposable { refreshQuickPickItems(); quickPick.show(); - quickPick.onDidAccept(async () => { - const selectedItem = quickPick.selectedItems[0]; - if ("cluster" in selectedItem) { - const cluster = selectedItem.cluster; - await this.connectionManager.attachCluster(cluster.id); - } else if (selectedItem.label === "$(cloud) Serverless") { - // Dispose the compute QuickPick before opening the version - // sub-picker so they don't stack visually. - disposables.forEach((d) => d.dispose()); - await this.selectServerless(); - return; - } else { - await UrlUtils.openExternal( - `${ - ( - await this.connectionManager.workspaceClient - ?.apiClient?.host - )?.href ?? "" - }#create/cluster` - ); - } - disposables.forEach((d) => d.dispose()); - }); + // `settled` guards a repeated Enter and stops the hide handler + // (which also fires when accept disposes the picker) from resolving + // before the accept branch's awaits finish. + return await new Promise((resolve) => { + let settled = false; + quickPick.onDidAccept(async () => { + if (settled) { + return; + } + settled = true; + let selected: SelectedCompute | undefined; + try { + const selectedItem = quickPick.selectedItems[0]; + if (selectedItem === undefined) { + // Accepted with nothing highlighted. + } else if ("cluster" in selectedItem) { + const cluster = selectedItem.cluster; + // Best-effort: attachCluster is @onError(throw:false), + // so a failed write pops its own error but doesn't + // throw. We still return the chosen target -- it's all + // setup-local needs and any failure is already shown. + await this.connectionManager.attachCluster( + cluster.id + ); + selected = {kind: "cluster", clusterId: cluster.id}; + } else if ( + selectedItem.label === "$(cloud) Serverless" + ) { + // Dispose the compute QuickPick before opening the + // version sub-picker so they don't stack visually. + disposables.forEach((d) => d.dispose()); + const version = await this.selectServerless(); + if (version !== undefined) { + selected = {kind: "serverless", version}; + } + } else { + await UrlUtils.openExternal( + `${ + ( + await this.connectionManager + .workspaceClient?.apiClient?.host + )?.href ?? "" + }#create/cluster` + ); + } + } catch (e) { + // Defense-in-depth for an unexpected throw (attach/enable + // are @onError and don't throw): keep it off the unhandled + // path and still settle. `selected` stays undefined. + NamedLogger.getOrCreate(Loggers.Extension).error( + "Compute picker selection failed", + e + ); + } finally { + disposables.forEach((d) => d.dispose()); + resolve(selected); + } + }); - quickPick.onDidHide(() => { - disposables.forEach((d) => d.dispose()); - quickPick.dispose(); + quickPick.onDidHide(() => { + disposables.forEach((d) => d.dispose()); + quickPick.dispose(); + if (!settled) { + settled = true; + resolve(undefined); + } + }); }); }; } @@ -242,18 +300,22 @@ export class ConnectionCommands implements Disposable { * selection, so setup need not re-prompt. If they dismiss the version * picker, no compute change is made. With the feature off this is the * plain, unchanged serverless enable. + * + * Returns the confirmed version, or `undefined` if the picker was dismissed + * or the feature is off (serverless enabled but version-less). */ - private async selectServerless() { + private async selectServerless(): Promise { if (!isPythonSetupEnabled()) { await this.connectionManager.enableServerless(); - return; + return undefined; } const version = await this.pickServerlessVersion(); if (version === undefined) { // User dismissed the version picker -- don't switch compute. - return; + return undefined; } await this.connectionManager.enableServerless(version); + return version; } /** diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index c34f05d1a..906d0161e 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -58,6 +58,7 @@ import { resolveComputeFrom, } from "./python-setup/controllers/pythonSetupDeps"; import {PythonSetupDriftManager} from "./python-setup/controllers/PythonSetupDriftManager"; +import {SetupCompute} from "./python-setup/controllers/PythonSetupEnvironmentSetup"; import {resolveCliPath} from "./python-setup/utils/setupLocalArgs"; import { isPythonSetupEnabled, @@ -1005,6 +1006,14 @@ export async function activate( // setup flow. persistServerlessVersion: (version) => connectionManager.enableServerless(version), + // Reuses the compute picker, which returns the chosen compute (or + // undefined if dismissed); its serverless branch is version-complete. + promptSelectCompute: () => + Promise.resolve( + commands.executeCommand( + "databricks.connection.attachClusterQuickPick" + ) + ), setActiveInterpreter: async (interpreterPath, root) => { await pythonExtensionWrapper.api.environments.updateActiveEnvironmentPath( interpreterPath, 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 abc4eeb4d..2386b0736 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -7,6 +7,7 @@ import { resolveComputeFrom, } from "./pythonSetupDeps"; import {PythonSetupState} from "../../vscode-objs/StateStorage"; +import {SetupCompute} from "./PythonSetupEnvironmentSetup"; import {Telemetry} from "../../telemetry"; import { SUCCESS_DEFAULT, @@ -175,6 +176,7 @@ function makeWiring( }), promptServerlessVersion: async () => "4", persistServerlessVersion: async () => {}, + promptSelectCompute: async () => undefined, setActiveInterpreter: async () => {}, persistSetupState: () => {}, log: {append: () => {}, show: () => {}}, @@ -254,21 +256,99 @@ describe("makePythonSetupDeps resolveCompute", () => { expect(prompted).to.equal(0); }); - it("never prompts when nothing is attached", async () => { - // Nothing selected is a real dead end, not a missing detail: the flow - // must guide the user, not open a version picker. - let prompted = 0; + it("offers the compute picker (not the version picker) when nothing is attached", async () => { + // Nothing attached opens the compute picker, never the version picker. + let versionPrompted = 0; + let pickerOpened = 0; const deps = makePythonSetupDeps( makeWiring({ + promptSelectCompute: async () => { + pickerOpened++; + return undefined; // dismissed + }, promptServerlessVersion: async () => { - prompted++; + versionPrompted++; return "4"; }, }) ); + // Dismissed -> `none`, which the orchestrator turns into guidance. expect(await deps.resolveCompute()).to.deep.equal({status: "none"}); - expect(prompted).to.equal(0); + expect(pickerOpened).to.equal(1); + expect(versionPrompted).to.equal(0); + }); + + it("runs against the cluster the picker returns, without re-reading state", async () => { + // Uses the picker's return value; `attachedCompute` stays `none` to + // prove the flow doesn't re-read (which would race the attach). + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => ({ + serverless: false, + cluster: undefined, + serverlessVersion: undefined, + }), + promptSelectCompute: async (): Promise< + SetupCompute | undefined + > => ({kind: "cluster", clusterId: "c1"}), + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "cluster", clusterId: "c1"}, + }); + }); + + it("runs against the serverless compute the picker returns", async () => { + // The picker returns serverless version-complete -- no follow-up prompt. + let versionPrompted = 0; + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => ({ + serverless: false, + cluster: undefined, + serverlessVersion: undefined, + }), + promptSelectCompute: async (): Promise< + SetupCompute | undefined + > => ({kind: "serverless", version: "5"}), + promptServerlessVersion: async () => { + versionPrompted++; + return "4"; + }, + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "serverless", version: "5"}, + }); + expect(versionPrompted).to.equal(0); + }); + + it("does not open the compute picker when a cluster is already attached", async () => { + let pickerOpened = 0; + const deps = makePythonSetupDeps( + makeWiring({ + attachedCompute: () => ({ + serverless: false, + cluster: {id: "c1"}, + serverlessVersion: undefined, + }), + promptSelectCompute: async () => { + pickerOpened++; + return undefined; + }, + }) + ); + + expect(await deps.resolveCompute()).to.deep.equal({ + status: "ok", + compute: {kind: "cluster", clusterId: "c1"}, + }); + expect(pickerOpened).to.equal(0); }); it("never prompts when serverless already has a version", async () => { diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index 2e294ab39..7de8b7b94 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -137,6 +137,13 @@ export interface PythonSetupWiringDeps { * selection, so the next run does not ask again. */ persistServerlessVersion: (version: string) => Promise; + /** + * Open the compute picker when nothing is attached, resolving to the chosen + * compute (or `undefined` if dismissed). The caller uses this return value + * directly, not a re-read of {@link attachedCompute}: a cluster attach + * propagates asynchronously, so an immediate re-read would race it. + */ + promptSelectCompute: () => Promise; /** 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. */ @@ -170,6 +177,15 @@ export function makePythonSetupDeps( isVisible, resolveCompute: async () => { const resolution = resolveComputeFrom(wiring.attachedCompute()); + if (resolution.status === "none") { + // Nothing attached: offer the picker inline instead of + // dead-ending. Use its return value (not a racy re-read); + // `undefined` means dismissed, so fall through to `none`. + const picked = await wiring.promptSelectCompute(); + return picked === undefined + ? {status: "none"} + : {status: "ok", compute: picked}; + } if (resolution.status !== "needsServerlessVersion") { return resolution; }