Start ssh tunnel

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.
 
-## <a id="select-cluster"></a>Select a cluster
+## <a id="select-cluster"></a>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.
 
 ## <a id="running-code"></a>Run Python code
 

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",

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);

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,

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;

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");
+    });
 });

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;

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();
         }
     }

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(