diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index 688416493..eac550bbe 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -48,13 +48,13 @@ The Databricks extension for Visual Studio Code enables you to connect to your r If your folder has multiple [Declarative Automation Bundles](#dabs), you can select which one to use by clicking "Open Existing Databricks project" button and selecting the desired project. -## Select a cluster +## Select compute -The extension uses an interactive cluster to run code. To select an interactive cluster: +The extension uses interactive compute to run code. To select interactive compute: 1. Open the Databricks panel by clicking on the Databricks icon on the left -2. Click on the "Select Cluster" button. - - If you wish to change the selected cluster, click on the "Configure Cluster" gear icon, next to the name of the selected cluster. +2. Click on the "Select Compute" button. + - If you wish to change the selected compute, click on the "Configure compute" gear icon, next to the name of the selected compute. ## Run Python code diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 651bf699e..10e248fcc 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -52,6 +52,13 @@ "types": "out/extension.d.ts", "contributes": { "commands": [ + { + "command": "databricks.ssh.startTunnel", + "title": "Start SSH Tunnel", + "category": "Databricks", + "icon": "$(remote)", + "enablement": "!databricks.context.remoteMode" + }, { "command": "databricks.connection.logout", "title": "Logout", @@ -73,20 +80,20 @@ }, { "command": "databricks.connection.attachCluster", - "title": "Attach cluster", + "title": "Attach compute", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(plug)" }, { "command": "databricks.connection.attachClusterQuickPick", - "title": "Configure cluster", + "title": "Configure compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(gear)" }, { "command": "databricks.connection.detachCluster", - "title": "Detach cluster", + "title": "Detach compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(debug-disconnect)" @@ -151,14 +158,14 @@ }, { "command": "databricks.cluster.start", - "title": "Start Cluster", + "title": "Start Compute", "icon": "$(debug-start)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" }, { "command": "databricks.cluster.stop", - "title": "Stop Cluster", + "title": "Stop Compute", "icon": "$(stop-circle)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" @@ -539,10 +546,15 @@ "name": "Configuration", "when": "!databricks.context.remoteMode" }, + { + "id": "sshTunnelView", + "name": "SSH Tunnel", + "when": "!databricks.context.remoteMode" + }, { "id": "clusterView", "when": "databricks.feature.views.cluster && !databricks.context.remoteMode", - "name": "Clusters" + "name": "Compute" }, { "id": "dabsResourceExplorerView", @@ -707,6 +719,10 @@ { "view": "configurationView", "contents": "To learn more about how to use the Databricks extension for Visual Studio Code [read our docs](https://docs.databricks.com/dev-tools/vscode-ext.html) or [Quickstart guide](command:databricks.quickstart.open)" + }, + { + "view": "sshTunnelView", + "contents": "Connect your IDE to a Databricks SSH Tunnel so you can run all Python and SQL workloads using Databricks compute.\n[Start SSH Tunnel](command:databricks.ssh.startTunnel)\nTo learn more about the SSH tunnel [read our docs](https://docs.databricks.com/aws/en/dev-tools/ssh-tunnel)." } ], "menus": { @@ -1233,7 +1249,7 @@ "submenus": [ { "id": "databricks.cluster.filter", - "label": "Filter clusters ...", + "label": "Filter compute ...", "icon": "$(filter)" }, { @@ -1439,7 +1455,7 @@ "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": false, - "description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)" + "description": "Enable/disable filtering for only accessible compute (compute on which the current user can run code)" }, "databricks.overrideDatabricksConfigFile": { "type": "string", diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index bd440c1d2..226bd0b08 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -6,7 +6,7 @@ import {execFile as execFileCb} from "node:child_process"; import {withFile} from "tmp-promise"; import {writeFile, readFile} from "node:fs/promises"; import {when, spy, reset, instance, mock} from "ts-mockito"; -import {CliWrapper, waitForProcess} from "./CliWrapper"; +import {CliWrapper, getSshConnectCommand, waitForProcess} from "./CliWrapper"; import path from "node:path"; import os from "node:os"; import crypto from "node:crypto"; @@ -128,6 +128,45 @@ describe(__filename, function () { assert.equal([command, ...args].join(" "), syncCommand); }); + it("should create ssh connect commands", () => { + // Logging is configured via env vars, not CLI flags, so no --log-* + // args appear on the ssh connect command line. + + // Serverless: no --cluster / --auto-start-cluster. + let {args} = getSshConnectCommand({compute: {type: "serverless"}}); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + ]); + + // Serverless GPU: --accelerator, no --cluster / --auto-start-cluster. + ({args} = getSshConnectCommand({ + compute: {type: "serverless", accelerator: "GPU_1xA10"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + "--accelerator=GPU_1xA10", + ]); + + // Dedicated cluster: --cluster and --auto-start-cluster. + ({args} = getSshConnectCommand({ + compute: {type: "cluster", clusterId: "1234-clusterid"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + "--cluster=1234-clusterid", + "--auto-start-cluster", + ]); + }); + it("should list profiles when no config file exists", async () => { const logFilePath = getTempLogFilePath(); const cli = createCliWrapper(logFilePath); diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 17ca9170b..77df0e7c9 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -10,6 +10,7 @@ import { Uri, commands, CancellationToken, + env, } from "vscode"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {promisify} from "node:util"; @@ -104,6 +105,37 @@ export interface ConfigEntry { } export type SyncType = "full" | "incremental"; + +export type SshConnectCompute = + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + +/** + * Constructs the `databricks ssh connect` command args for opening a remote + * IDE window. Serverless is the default when no cluster is given. + * + * The --ide flag matches the host editor so the CLI opens the right remote + * window: Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. + * + * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see + * CliWrapper.getSshConnectEnvVars), so we do not pass --log-* flags here. + */ +export function getSshConnectCommand(opts: {compute: SshConnectCompute}): { + args: string[]; +} { + const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; + const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; + if (opts.compute.type === "cluster") { + // Start a stopped single-user cluster when connecting. + args.push(`--cluster=${opts.compute.clusterId}`); + args.push("--auto-start-cluster"); + } else if (opts.compute.accelerator) { + // Serverless GPU: request a specific accelerator type. + args.push(`--accelerator=${opts.compute.accelerator}`); + } + return {args}; +} export class ProcessError extends Error { constructor( message: string, @@ -574,6 +606,24 @@ export class CliWrapper { }); } + /** + * Env vars for interactive CLI commands run in a terminal (e.g. `ssh + * connect`). Auth is forwarded via env vars, matching the bundle init flow. + */ + getSshConnectEnvVars(authProvider: AuthProvider) { + return removeUndefinedKeys({ + ...EnvVarGenerators.getEnvVarsForCli( + this.extensionContext, + workspaceConfigs.databrickscfgLocation + ), + ...EnvVarGenerators.getProxyEnvVars(), + ...this.getLogginEnvVars(), + ...authProvider.toEnv(), + // eslint-disable-next-line @typescript-eslint/naming-convention + DATABRICKS_OUTPUT_FORMAT: "text", + }); + } + async bundleInit( templateDirPath: string, outputDirPath: string, diff --git a/packages/databricks-vscode/src/cluster/ClusterModel.ts b/packages/databricks-vscode/src/cluster/ClusterModel.ts index b1ab85dd2..a16abcb03 100644 --- a/packages/databricks-vscode/src/cluster/ClusterModel.ts +++ b/packages/databricks-vscode/src/cluster/ClusterModel.ts @@ -67,6 +67,16 @@ export class ClusterModel implements Disposable { ); } + /** + * All loaded clusters, ignoring the explorer's active filter. Consumers that + * need the complete set (e.g. the SSH tunnel picker, which starts stopped + * clusters via --auto-start-cluster) must not inherit the explorer's + * `ALL`/`ME`/`RUNNING` filter. + */ + public get allRoots(): Cluster[] | undefined { + return sortClusters(Array.from(this.clusterLoader.clusters.values())); + } + private applyFilter(nodes: Cluster[] | undefined): Cluster[] | undefined { if (!nodes) { return nodes; diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts index b0facc8ba..626557093 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts @@ -3,7 +3,10 @@ import {ApiClient} from "@databricks/sdk-experimental"; import {Cluster} from "../sdk-extensions"; import assert from "assert"; import {mock} from "ts-mockito"; -import {formatQuickPickClusterDetails} from "./ConnectionCommands"; +import { + formatClusterState, + formatQuickPickClusterDetails, +} from "./ConnectionCommands"; describe(__filename, () => { it("attach cluster quickpick: correctly format cluster details", () => { @@ -21,4 +24,14 @@ describe(__filename, () => { assert.equal(clusterDetails, `2 GB | 4 Cores | spark-version | user-2`); }); + + it("formatClusterState: maps RUNNING/TERMINATED to Active/Inactive and title-cases the rest", () => { + assert.equal(formatClusterState("RUNNING"), "Active"); + assert.equal(formatClusterState("TERMINATED"), "Inactive"); + assert.equal(formatClusterState("PENDING"), "Pending"); + assert.equal(formatClusterState("RESTARTING"), "Restarting"); + assert.equal(formatClusterState("TERMINATING"), "Terminating"); + assert.equal(formatClusterState("ERROR"), "Error"); + assert.equal(formatClusterState("UNKNOWN"), "Unknown"); + }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index abea24dff..091892199 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -1,4 +1,5 @@ import {Cluster} from "../sdk-extensions"; +import {compute} from "@databricks/sdk-experimental"; import { Disposable, QuickPickItem, @@ -40,6 +41,20 @@ function formatQuickPickClusterSize(sizeInMB: number): string { return `${sizeInMB} MB`; } } +// Formats a compute state for display in the picker. RUNNING/TERMINATED map to +// the softer "Active"/"Inactive" (matching the Workspace UI); other states are +// title-cased ("PENDING" -> "Pending") so they read less harshly. +export function formatClusterState(state: compute.State): string { + switch (state) { + case "RUNNING": + return "Active"; + case "TERMINATED": + return "Inactive"; + default: + return state.charAt(0) + state.slice(1).toLowerCase(); + } +} + export function formatQuickPickClusterDetails(cluster: Cluster) { const details = []; if (cluster.memoryMb) { @@ -143,7 +158,7 @@ export class ConnectionCommands implements Disposable { ClusterItem | QuickPickItem >(); quickPick.title = - typeof title === "string" ? title : "Select Cluster"; + typeof title === "string" ? title : "Select Compute"; quickPick.keepScrollPosition = true; quickPick.busy = true; quickPick.canSelectMany = false; diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 50b7e6999..37985f6ac 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -68,6 +68,19 @@ export class ConnectionManager implements Disposable { this.onDidChangeSyncDestinationEmitter.event; private readonly initialization = new Barrier(); + // Set once init() has resolved the initialization barrier. Guards callers + // (e.g. the SSH tunnel flow) that must not `await login()` before init(), + // since that would block on the barrier forever when the workspace is not a + // Databricks project and init() is never called. + private _initialized = false; + + /** + * Whether the connection manager has been initialized (init() ran). + * Only true once the workspace has been set up as a Databricks project. + */ + get isInitialized(): boolean { + return this._initialized; + } get projectRoot() { return this.workspaceFolderManager.activeProjectUri; @@ -228,6 +241,7 @@ export class ConnectionManager implements Disposable { ) ) ); + this._initialized = true; this.initialization.resolve(); } } diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index b06c32dc1..9d9e00816 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -53,6 +53,7 @@ import {Events, Metadata} from "./telemetry/constants"; import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller"; import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits"; import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton"; +import {SshCommands} from "./ssh/SshCommands"; import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager"; import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue"; import { @@ -169,6 +170,18 @@ export async function activate( } ) ); + // The SSH Tunnel panel is always visible, including on the start screen + // with no folder open. There is no ConnectionManager/ClusterModel here, + // so wire the command to the standalone (login-then-tunnel) flow. + const sshCommands = new SshCommands(cli); + context.subscriptions.push( + sshCommands, + telemetry.registerCommand( + "databricks.ssh.startTunnel", + sshCommands.startTunnelCommand, + sshCommands + ) + ); // We show a welcome view when there's no workspace folders, prompting users // to either open a new folder or to initialize a new databricks project. // In both cases we expect the workspace to be reloaded and the extension will @@ -836,6 +849,17 @@ export async function activate( ) ); + // SSH tunnel (remote development) group + const sshCommands = new SshCommands(cli, connectionManager, clusterModel); + context.subscriptions.push( + sshCommands, + telemetry.registerCommand( + "databricks.ssh.startTunnel", + sshCommands.startTunnelCommand, + sshCommands + ) + ); + // Cluster group const clusterTreeDataProvider = new ClusterListDataProvider(clusterModel); const clusterCommands = new ClusterCommands( diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts new file mode 100644 index 000000000..2ea734a62 --- /dev/null +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -0,0 +1,441 @@ +import { + Disposable, + Event, + EventEmitter, + QuickPick, + QuickPickItem, + QuickPickItemKind, + window, +} from "vscode"; +import {WorkspaceClient} from "@databricks/sdk-experimental"; +import {ClusterModel} from "../cluster/ClusterModel"; +import {ConnectionManager} from "../configuration/ConnectionManager"; +import { + ClusterItem, + formatClusterState, + formatQuickPickClusterDetails, +} from "../configuration/ConnectionCommands"; +import { + CliWrapper, + getSshConnectCommand, + SshConnectCompute, +} from "../cli/CliWrapper"; +import {AuthProvider} from "../configuration/auth/AuthProvider"; +import {LoginWizard} from "../configuration/LoginWizard"; +import {Cluster} from "../sdk-extensions"; +import {onError} from "../utils/onErrorDecorator"; +import {logging} from "@databricks/sdk-experimental"; +import {Loggers} from "../logger"; + +const SERVERLESS_LABEL = "$(cloud) Serverless"; + +type Compute = SshConnectCompute; + +/** + * A serverless QuickPick item, optionally carrying a GPU accelerator type + * forwarded to `databricks ssh connect --accelerator`. Plain serverless has no + * accelerator. + */ +interface ServerlessItem extends QuickPickItem { + accelerator?: string; +} + +// Serverless compute options shown at the top of the picker: plain serverless +// plus the serverless GPU accelerator types supported by the CLI. +const SERVERLESS_ITEMS: ServerlessItem[] = [ + { + label: SERVERLESS_LABEL, + alwaysShow: true, + }, + { + label: "$(cloud) Serverless GPU 1xA10", + alwaysShow: true, + accelerator: "GPU_1xA10", + }, + { + label: "$(cloud) Serverless GPU 8xH100", + alwaysShow: true, + accelerator: "GPU_8xH100", + }, +]; + +/** + * Minimal cluster feed the compute picker needs. `ClusterModel` (connected + * path) satisfies this directly; `StandaloneClusterSource` (start-screen path, + * no workspace folder) implements the same shape from a bare WorkspaceClient. + */ +interface ClusterSource extends Disposable { + readonly roots: Cluster[] | undefined; + readonly onDidChange: Event; + refresh(): void; +} + +/** + * Fetches eligible clusters directly from a WorkspaceClient for the standalone + * (no workspace folder) tunnel flow, where no `ClusterModel` exists. + */ +class StandaloneClusterSource implements ClusterSource { + private _clusters: Cluster[] | undefined; + private readonly onDidChangeEmitter = new EventEmitter(); + readonly onDidChange = this.onDidChangeEmitter.event; + + constructor(private readonly workspaceClient: WorkspaceClient) {} + + get roots(): Cluster[] | undefined { + return this._clusters; + } + + refresh() { + void this.load(); + } + + private async load() { + try { + const clusters: Cluster[] = []; + for await (const cluster of Cluster.list( + this.workspaceClient.apiClient + )) { + clusters.push(cluster); + } + this._clusters = clusters; + } catch (e) { + // On a list failure (network, permissions, expired token) surface an + // empty set rather than leaving the picker spinning until the 10s + // fallback, and still fire onDidChange so the state is honest. + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Failed to list clusters for SSH tunnel", + e + ); + this._clusters = []; + } + this.onDidChangeEmitter.fire(); + } + + dispose() { + this.onDidChangeEmitter.dispose(); + } +} + +/** + * Adapts the shared `ClusterModel` (connected path) to `ClusterSource`, exposing + * its *unfiltered* cluster set so the picker doesn't inherit the explorer's + * `ALL`/`ME`/`RUNNING` filter. The model is owned by the extension, so `dispose` + * is a no-op here. + */ +class ClusterModelSource implements ClusterSource { + constructor(private readonly clusterModel: ClusterModel) {} + + get roots(): Cluster[] | undefined { + return this.clusterModel.allRoots; + } + + get onDidChange(): Event { + return this.clusterModel.onDidChange; + } + + refresh() { + this.clusterModel.refresh(); + } + + dispose() {} +} + +/** + * The auth + compute context needed to launch a tunnel, resolved either from an + * already-connected workspace or from a standalone login on the start screen. + */ +interface TunnelContext { + authProvider: AuthProvider; + userName: string; + clusterSource: ClusterSource; + // Cluster sources we create ourselves (standalone) must be disposed after + // the picker; the shared ClusterModel is owned by the extension and is not. + ownsClusterSource: boolean; +} + +export class SshCommands implements Disposable { + private disposables: Disposable[] = []; + + /** + * `connectionManager`/`clusterModel` are only available once a workspace + * folder is open. When they are undefined (start screen) the command falls + * back to a standalone login flow so the tunnel can be started from the + * dedicated SSH Tunnel panel with no folder open. + */ + constructor( + private readonly cli: CliWrapper, + private readonly connectionManager?: ConnectionManager, + private readonly clusterModel?: ClusterModel + ) {} + + @onError({popup: {prefix: "Error starting SSH tunnel."}}) + async startTunnelCommand() { + const context = await this.resolveTunnelContext(); + if (context === undefined) { + return; + } + try { + const compute = await this.pickCompute( + context.userName, + context.clusterSource + ); + if (compute === undefined) { + return; + } + await this.launchSshTunnel(context.authProvider, compute); + } finally { + if (context.ownsClusterSource) { + context.clusterSource.dispose(); + } + } + } + + /** + * Resolves the auth provider, user and cluster feed for the tunnel. Uses the + * connected workspace when available (connecting first if needed), otherwise + * runs a standalone login wizard so the tunnel works with no folder open. + */ + private async resolveTunnelContext(): Promise { + // Only use the connected path when the connection manager has actually + // been initialized. When a folder is open but is not a Databricks + // project, init() is never called, so awaiting login() would block on + // its initialization barrier forever (the button would silently do + // nothing). In that case we fall through to the standalone login flow. + if (this.connectionManager?.isInitialized && this.clusterModel) { + if (this.connectionManager.state !== "CONNECTED") { + await this.connectionManager.login(true); + } + const workspace = this.connectionManager.databricksWorkspace; + if (!workspace || this.connectionManager.state !== "CONNECTED") { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return undefined; + } + return { + authProvider: workspace.authProvider, + userName: workspace.userName, + // Read the unfiltered cluster set so the picker is independent of + // the explorer's active filter — a "Running" filter would + // otherwise hide stopped single-user clusters that this flow can + // still start via --auto-start-cluster. + clusterSource: new ClusterModelSource(this.clusterModel), + ownsClusterSource: false, + }; + } + + const authProvider = await LoginWizard.run(this.cli); + if (authProvider === undefined || !(await authProvider.check())) { + return undefined; + } + const workspaceClient = await authProvider.getWorkspaceClient(); + const userName = (await workspaceClient.currentUser.me()).userName; + if (!userName) { + window.showErrorMessage( + "Could not determine the current user for the SSH tunnel." + ); + return undefined; + } + return { + authProvider, + userName, + clusterSource: new StandaloneClusterSource(workspaceClient), + ownsClusterSource: true, + }; + } + + private pickCompute( + me: string, + clusterSource: ClusterSource + ): Promise { + return new Promise((resolve) => { + const quickPick = window.createQuickPick< + ClusterItem | ServerlessItem + >(); + quickPick.title = "Select compute for SSH tunnel"; + quickPick.placeholder = "Loading dedicated clusters…"; + quickPick.keepScrollPosition = true; + quickPick.busy = true; + quickPick.canSelectMany = false; + + // Whether the dedicated-cluster list is still being fetched. Drives + // the separator label and placeholder so the user can tell the list + // is still loading rather than waiting on them to pick. + let loading = true; + + // Preselect only once, on the first populated repaint. refreshItems + // runs on every onDidChange (the loader fires once per cluster + // permission check), and re-applying activeItems would keep yanking + // the highlight back to the preselected entry while the user + // navigates. + let hasPreselected = false; + + // Serverless items are always ready; the separator that follows them + // reflects the dedicated-cluster loading state. + const buildStaticItems = ( + clusterCount: number + ): (ServerlessItem | QuickPickItem)[] => [ + ...SERVERLESS_ITEMS, + { + label: loading + ? "Loading clusters…" + : clusterCount > 0 + ? "Dedicated clusters" + : "No dedicated clusters found", + kind: QuickPickItemKind.Separator, + }, + ]; + + const stopLoading = () => { + loading = false; + quickPick.busy = false; + quickPick.placeholder = undefined; + }; + + quickPick.items = buildStaticItems(0); + + const refreshItems = () => { + // Only dedicated single-user clusters owned by the current user + // can be used for an SSH tunnel. + const clusters = (clusterSource.roots ?? []).filter((c) => + c.isValidSingleUser(me) + ); + // Clear the loading state only once clusters have actually + // loaded, before building items so the separator label matches + // this repaint. On a cold first open the loader is still + // fetching, so we keep loading and let onDidChange repaint when + // clusters arrive. + if (clusters.length > 0) { + stopLoading(); + } + quickPick.items = buildStaticItems(clusters.length).concat( + clusters.map((c) => ({ + // Use the same compute icon as the "Select a compute" + // entry in the configuration section; the per-state icon + // is dropped in favour of showing the state as text. + label: `$(server) ${c.name!} (${c.id})`, + detail: `${formatClusterState( + c.state + )} | ${formatQuickPickClusterDetails(c)}`, + cluster: c, + })) + ); + if (!hasPreselected) { + hasPreselected = this.preselect(quickPick); + } + }; + + // Fallback so the spinner can't hang forever for a user with no + // eligible clusters (onDidChange may never add any). Repaint so the + // separator flips to "No dedicated clusters found". + const spinnerTimeout = setTimeout(() => { + stopLoading(); + refreshItems(); + }, 10_000); + + // Register the change listener before triggering refresh() so no + // onDidChange fired by the (re)started loader can be missed. + const disposables: Disposable[] = [ + clusterSource.onDidChange(refreshItems), + quickPick, + {dispose: () => clearTimeout(spinnerTimeout)}, + ]; + + // Paint whatever is already cached first (fast path on reopen), then + // trigger a reload; fresh results stream in via onDidChange. + refreshItems(); + clusterSource.refresh(); + quickPick.show(); + + quickPick.onDidAccept(() => { + const selectedItem = quickPick.selectedItems[0]; + // Resolve with the accepted value BEFORE disposing. Disposing + // the quickPick (it is in `disposables`) synchronously fires + // onDidHide, whose handler resolves undefined; since resolve is + // idempotent, settling the real value first makes that later + // undefined a harmless no-op. Disposing first would let + // undefined win and drop the user's selection. + if (selectedItem === undefined) { + resolve(undefined); + } else if ("cluster" in selectedItem) { + resolve({ + type: "cluster", + clusterId: selectedItem.cluster.id, + }); + } else { + resolve({ + type: "serverless", + accelerator: selectedItem.accelerator, + }); + } + disposables.forEach((d) => d.dispose()); + }); + + quickPick.onDidHide(() => { + disposables.forEach((d) => d.dispose()); + resolve(undefined); + }); + }); + } + + /** + * Pre-selects the compute the user already has configured locally: the + * attached single-user cluster if any, otherwise serverless. Only the + * connected path has a configured cluster/serverless preference. Returns + * whether a selection was actually applied so the caller can stop + * re-preselecting once the target item exists (the attached cluster may not + * be in the list yet on a cold first repaint). + */ + private preselect( + quickPick: QuickPick + ): boolean { + const currentCluster = this.connectionManager?.cluster; + if (currentCluster?.isSingleUser()) { + const match = quickPick.items.find( + (i): i is ClusterItem => + "cluster" in i && i.cluster.id === currentCluster.id + ); + if (match) { + quickPick.activeItems = [match]; + return true; + } + } + if (this.connectionManager?.serverless) { + const serverlessItem = quickPick.items.find( + (i) => i.label === SERVERLESS_LABEL + ); + if (serverlessItem) { + quickPick.activeItems = [serverlessItem]; + return true; + } + } + return false; + } + + private async launchSshTunnel( + authProvider: AuthProvider, + compute: Compute + ) { + const {args} = getSshConnectCommand({compute}); + + const env: Record = { + ...this.cli.getSshConnectEnvVars(authProvider), + }; + + // The transient terminal is not retained on `this.disposables`: doing so + // would only release it on extension deactivate/reload, tearing down an + // active tunnel. It is `isTransient` and disposes itself when closed. + const terminal = window.createTerminal({ + name: "Databricks SSH Tunnel", + isTransient: true, + env, + strictEnv: false, + }); + terminal.show(); + terminal.sendText(`${this.cli.escapedCliPath} ${args.join(" ")}`); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } +} diff --git a/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts b/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts index 63379ba2f..ce3057cc9 100644 --- a/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/run_dbconnect.ucws.e2e.ts @@ -643,7 +643,7 @@ describe("Run files on serverless compute", async function () { }); it("should select serverless compute", async () => { - await executeCommandWhenAvailable("Databricks: Configure cluster"); + await executeCommandWhenAvailable("Databricks: Configure compute"); const computeInput = await waitForInput(); await computeInput.selectQuickPick("Serverless"); }); diff --git a/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts b/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts index c97b89d7c..71ac38e3b 100644 --- a/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts +++ b/packages/databricks-vscode/src/test/e2e/utils/commonUtils.ts @@ -12,6 +12,7 @@ import { const ViewSectionTypes = [ "CLUSTERS", "CONFIGURATION", + "SSH TUNNEL", "WORKSPACE FILE SYSTEM", "BUNDLE RESOURCE EXPLORER", "BUNDLE VARIABLES", diff --git a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts index 6a0814f74..4c216b529 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts @@ -114,7 +114,7 @@ export class ClusterComponent extends BaseComponent { // We are logged in -> Select cluster prompt return [ { - label: LabelUtils.highlightedLabel("Select a cluster"), + label: LabelUtils.highlightedLabel("Select compute"), collapsibleState: TreeItemCollapsibleState.None, contextValue: getContextValue("none"), iconPath: new ThemeIcon( @@ -123,7 +123,7 @@ export class ClusterComponent extends BaseComponent { ), id: TREE_ICON_ID, command: { - title: "Select a cluster", + title: "Select compute", command: "databricks.connection.attachClusterQuickPick", }, },