[SSH] Report port forwarding refusals as their own telemetry categories

.nextchanges/cli/ssh-port-forwarding-disabled.md

@@ -0,0 +1 @@
+* `databricks ssh connect` now fails right away with the workspace's reason when port forwarding to serverless compute is disabled, instead of waiting out the server start timeout. ([#7057](https://github.com/databricks/cli/pull/7057))

acceptance/ssh/connect-serverless-feature-disabled/out.connect.txt

@@ -0,0 +1,5 @@
+Connecting to feature-disabled...
+Uploading binaries...
+Starting SSH server...
+Job submitted successfully with run ID: [RUN_ID]
+Error: failed to ensure that ssh server is running: the workspace refused the SSH tunnel connection: FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace, which includes the SSH tunnel used by `databricks ssh connect`

acceptance/ssh/connect-serverless-feature-disabled/out.test.toml

@@ -0,0 +1,5 @@
+Cloud = false
+GOOS.darwin = false
+GOOS.linux = true
+GOOS.windows = false
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/ssh/connect-serverless-feature-disabled/output.txt

@@ -0,0 +1,18 @@
+
+>>> [CLI] workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json {"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}
+
+=== CLI ssh connect
+
+=== Reported error category
+
+>>> jq -r select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category [OUTPUT_DIR]/out.requests.txt
+DRIVER_PROXY_FORBIDDEN
+
+>>> print_requests.py //api/2.2/jobs/runs/cancel
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/runs/cancel",
+  "body": {
+    "run_id": [RUN_ID]
+  }
+}

acceptance/ssh/connect-serverless-feature-disabled/script

@@ -0,0 +1,21 @@
+# An administrator disabled port forwarding to serverless compute, so the driver proxy refuses
+# every request to the tunnel server. The setting itself cannot be read here (see test.toml), so
+# the CLI goes ahead and starts the server. It must then report the proxy's reason as soon as it
+# gets the 403 instead of polling until the server start timeout, and cancel the run it submitted
+# since nothing can ever connect to that server.
+# The body repeats the name: with only {"value": false} in it, the API sees an empty message.
+trace $CLI workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json '{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}' >/dev/null
+
+# No release artifacts locally and they're never executed; stand in empty archives.
+mkdir -p releases
+: >releases/databricks_cli_linux_amd64.zip
+: >releases/databricks_cli_linux_arm64.zip
+
+title "CLI ssh connect\n"
+musterr $CLI ssh connect --name feature-disabled --releases-dir="$PWD/releases" -- "echo 'Connection successful'" >out.connect.txt 2>&1
+add_repl "$(grep -o 'run ID: [0-9]*' out.connect.txt | grep -o '[0-9]*$')" RUN_ID
+
+title "Reported error category\n"
+trace jq -r 'select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category' "$OUT_REQUESTS"
+
+trace print_requests.py //api/2.2/jobs/runs/cancel

acceptance/ssh/connect-serverless-feature-disabled/test.toml

@@ -0,0 +1,29 @@
+Cloud = false
+
+# Record requests so we can assert that the CLI cancels the run it submitted.
+RecordRequests = true
+
+# Stand-in release archives staged by the script; not golden output.
+Ignore = [
+  "releases",
+]
+
+# The rest of acceptance/ssh only runs on Linux; nothing here needs another OS.
+GOOS.linux = true
+GOOS.darwin = false
+GOOS.windows = false
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+# Production does not serve port_fwd_serverless_enabled through the public API, so the CLI cannot
+# check it up front there and only learns about it from the driver proxy's 403. This reproduces
+# that environment, which the fake workspace's settings handler does not model.
+[[Server]]
+Pattern = "GET /api/2.1/settings/port_fwd_serverless_enabled"
+Response.StatusCode = 404
+Response.Body = '''
+{
+  "error_code": "NOT_FOUND",
+  "message": "Setting port_fwd_serverless_enabled is not enabled for public API"
+}
+'''

acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml

@@ -0,0 +1,5 @@
+Cloud = false
+GOOS.darwin = false
+GOOS.linux = true
+GOOS.windows = false
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt

@@ -0,0 +1,15 @@
+
+>>> [CLI] workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json {"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}
+
+=== CLI ssh connect
+
+>>> musterr [CLI] ssh connect --name port-forwarding-disabled -- echo 'Connection successful'
+Connecting to port-forwarding-disabled...
+Error: port forwarding to serverless compute is disabled for this workspace (setting port_fwd_serverless_enabled), and the SSH tunnel needs it. Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster
+
+=== Reported error category
+
+>>> jq -r select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category [OUTPUT_DIR]/out.requests.txt
+PORT_FORWARDING_DISABLED
+
+>>> print_requests.py //api

acceptance/ssh/connect-serverless-port-forwarding-disabled/script

@@ -0,0 +1,16 @@
+# An administrator disabled port forwarding to serverless compute. The CLI reads the setting
+# before anything else, so it refuses the connection without creating a secret scope, uploading
+# binaries or starting compute that nobody could reach.
+# The body repeats the name: with only {"value": false} in it, the API sees an empty message.
+trace $CLI workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json '{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}' >/dev/null
+# Drop the setup request so the check below only sees what ssh connect sent.
+print_requests.py //api/2.1/settings >/dev/null
+
+title "CLI ssh connect\n"
+trace musterr $CLI ssh connect --name port-forwarding-disabled -- "echo 'Connection successful'"
+
+title "Reported error category\n"
+trace jq -r 'select(.path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.ssh_tunnel_event.error_category' "$OUT_REQUESTS"
+
+# Expect no requests other than reads: the connection stops at the settings check.
+trace print_requests.py //api

acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml

@@ -0,0 +1,11 @@
+Cloud = false
+
+# Record requests so we can assert that nothing is created before the CLI gives up.
+RecordRequests = true
+
+# The rest of acceptance/ssh only runs on Linux; nothing here needs another OS.
+GOOS.linux = true
+GOOS.darwin = false
+GOOS.windows = false
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

experimental/ssh/internal/client/client.go

@@ -38,6 +38,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/environments"
 	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/gorilla/websocket"
 )
@@ -47,6 +48,19 @@ var sshServerBootstrapScript string
 
 var errServerMetadata = errors.New("server metadata error")
 
+// portFwdServerlessSetting is the workspace setting an administrator turns off to disable port
+// forwarding to serverless compute, which the tunnel needs.
+const portFwdServerlessSetting = "port_fwd_serverless_enabled"
+
+var errPortForwardingDisabled = errors.New("port forwarding to serverless compute is disabled for this workspace " +
+	"(setting " + portFwdServerlessSetting + "), and the SSH tunnel needs it. " +
+	"Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster")
+
+// errDriverProxyForbidden is a 403 from the driver proxy. Unlike the other non-200 responses, which a
+// server that is still starting produces too, retrying cannot change it: the workspace forbids the
+// connection, e.g. because an administrator disabled port forwarding to serverless compute.
+var errDriverProxyForbidden = errors.New("the workspace refused the SSH tunnel connection")
+
 var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)
 
 const (
@@ -341,6 +355,13 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID))
 	}
 
+	if opts.IsServerlessMode() {
+		if err := checkPortForwardingToServerless(ctx, client); err != nil {
+			outcome.errorCategory = protos.SshTunnelErrorCategoryPortForwardingDisabled
+			return err
+		}
+	}
+
 	if opts.IDE != "" && !opts.ProxyMode {
 		if err := vscode.CheckIDECommand(opts.IDE); err != nil {
 			outcome.errorCategory = protos.SshTunnelErrorCategoryIDECommandNotOnPath
@@ -445,6 +466,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		userName, serverPort, clusterID, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts)
 		if err != nil {
 			outcome.errorCategory = protos.SshTunnelErrorCategoryServerStartTimeout
+			if errors.Is(err, errDriverProxyForbidden) {
+				outcome.errorCategory = protos.SshTunnelErrorCategoryDriverProxyForbidden
+			}
 			return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
 		}
 		outcome.serverStartTimeMs = time.Since(serverStartTime).Milliseconds()
@@ -653,7 +677,7 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	log.Debugf(ctx, "Metadata response status code: %d", resp.StatusCode)
 
 	if resp.StatusCode != http.StatusOK {
-		return serverMetadata{}, errors.Join(errServerMetadata, fmt.Errorf("server is not ok, status code %d", resp.StatusCode))
+		return serverMetadata{}, metadataResponseError(resp.StatusCode, bodyBytes)
 	}
 
 	return serverMetadata{
@@ -665,6 +689,46 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	}, nil
 }
 
+// metadataResponseError classifies a non-200 response from the server's /metadata endpoint. A 403
+// is final and carries the driver proxy's own explanation; anything else is errServerMetadata,
+// which callers treat as a server that is not running yet.
+func metadataResponseError(statusCode int, body []byte) error {
+	if statusCode != http.StatusForbidden {
+		return errors.Join(errServerMetadata, fmt.Errorf("server is not ok, status code %d", statusCode))
+	}
+	// The driver proxy writes its error as a JSON string ("CODE: message") rather than the usual
+	// {"error_code", "message"} object.
+	var message string
+	if json.Unmarshal(body, &message) != nil {
+		message = string(body)
+	}
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = fmt.Sprintf("status code %d", statusCode)
+	}
+	return fmt.Errorf("%w: %s", errDriverProxyForbidden, message)
+}
+
+// checkPortForwardingToServerless refuses a serverless connection before any compute is started when
+// the workspace setting that the driver proxy enforces for the tunnel is off. The proxy's gate
+// (PortForwardingToServerlessGate in data-plane-router) also depends on an enforcement flag and a port
+// allow-list that the CLI cannot see, and the setting is not readable through the public API in every
+// environment. A read that fails therefore lets the connection go ahead, and the proxy's own 403
+// (errDriverProxyForbidden) remains the backstop.
+func checkPortForwardingToServerless(ctx context.Context, client *databricks.WorkspaceClient) error {
+	setting, err := client.WorkspaceSettingsV2.GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{
+		Name: portFwdServerlessSetting,
+	})
+	if err != nil {
+		log.Debugf(ctx, "Could not read the %s setting, continuing: %v", portFwdServerlessSetting, err)
+		return nil
+	}
+	if setting.EffectiveBooleanVal != nil && !setting.EffectiveBooleanVal.Value {
+		return errPortForwardingDisabled
+	}
+	return nil
+}
+
 // newDriverProxyRequest builds an authenticated GET request to one of the SSH server's
 // HTTP endpoints behind the workspace driver proxy.
 func newDriverProxyRequest(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, port int, endpoint, liteswap string) (*http.Request, error) {
@@ -1292,6 +1356,15 @@ func describeRunFailure(ctx context.Context, client *databricks.WorkspaceClient,
 	return strings.TrimRight(b.String(), "\n")
 }
 
+// cancelSSHServerRun cancels the bootstrap run this connection submitted once no client can reach
+// its server. Best effort: the server also shuts itself down after its no-client delay, so a
+// failure here must not hide the error that made the connection give up.
+func cancelSSHServerRun(ctx context.Context, client *databricks.WorkspaceClient, runID int64) {
+	if _, err := client.Jobs.CancelRun(ctx, jobs.CancelRun{RunId: runID}); err != nil {
+		log.Debugf(ctx, "Failed to cancel the ssh server job run %d: %v", runID, err)
+	}
+}
+
 // runFailureIfTerminated reports whether the bootstrap run has reached a terminal state (so the
 // SSH server will never come up), returning a formatted failure description when it has.
 func runFailureIfTerminated(ctx context.Context, client *databricks.WorkspaceClient, runID int64) (string, bool) {
@@ -1381,6 +1454,7 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 	clusterID := opts.ClusterID
 
 	meta, err := getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
+	// This includes errDriverProxyForbidden: a new server would be refused the same way.
 	if err != nil && !errors.Is(err, errServerMetadata) {
 		return "", 0, "", err
 	}
@@ -1410,6 +1484,10 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 				return "", 0, "", ctx.Err()
 			}
 			meta, err = getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
+			if errors.Is(err, errDriverProxyForbidden) {
+				cancelSSHServerRun(ctx, client, runID)
+				return "", 0, "", err
+			}
 			// Accept only once metadata reflects the requested usage policy, so we don't latch
 			// onto a server a previous connection started under a different policy before our new
 			// server has overwritten metadata.json.

experimental/ssh/internal/client/client_internal_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/environments"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -803,3 +804,97 @@ func TestBuildSshTunnelEventReportsFailure(t *testing.T) {
 	require.NoError(t, err)
 	assert.Contains(t, string(b), `"is_success":false`)
 }
+
+func TestMetadataResponseError(t *testing.T) {
+	tests := []struct {
+		name       string
+		statusCode int
+		body       string
+		forbidden  bool
+		message    string
+	}{
+		{
+			name:       "feature disabled",
+			statusCode: 403,
+			body:       `"FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace"`,
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace",
+		},
+		{
+			name:       "forbidden with a body that is not a JSON string",
+			statusCode: 403,
+			body:       "Forbidden\n",
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: Forbidden",
+		},
+		{
+			name:       "forbidden with an empty body",
+			statusCode: 403,
+			body:       "",
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: status code 403",
+		},
+		{
+			// A stale metadata.json points at a terminated cluster; the caller must start a new server.
+			name:       "terminated cluster is retried",
+			statusCode: 400,
+			body:       `"INVALID_STATE: Cluster 1014-135824-abcdefgh-v2n is in Terminated state but Spark driver is not healthy"`,
+			forbidden:  false,
+			message:    "server metadata error\nserver is not ok, status code 400",
+		},
+		{
+			name:       "server still starting is retried",
+			statusCode: 502,
+			body:       "",
+			forbidden:  false,
+			message:    "server metadata error\nserver is not ok, status code 502",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := metadataResponseError(tc.statusCode, []byte(tc.body))
+			assert.Equal(t, tc.forbidden, errors.Is(err, errDriverProxyForbidden))
+			assert.Equal(t, !tc.forbidden, errors.Is(err, errServerMetadata))
+			assert.EqualError(t, err, tc.message)
+		})
+	}
+}
+
+func TestCheckPortForwardingToServerless(t *testing.T) {
+	tests := []struct {
+		name    string
+		setting *settingsv2.Setting
+		readErr error
+		wantErr error
+	}{
+		{
+			name:    "disabled",
+			setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: false}},
+			wantErr: errPortForwardingDisabled,
+		},
+		{
+			name:    "enabled",
+			setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: true}},
+		},
+		{
+			name:    "no effective value",
+			setting: &settingsv2.Setting{},
+		},
+		{
+			// Production does not serve the setting through the public API yet.
+			name:    "read fails",
+			readErr: errors.New("Setting port_fwd_serverless_enabled is not enabled for public API"),
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := cmdio.MockDiscard(t.Context())
+			m := mocks.NewMockWorkspaceClient(t)
+			m.GetMockWorkspaceSettingsV2API().EXPECT().
+				GetPublicWorkspaceSetting(mock.Anything, settingsv2.GetPublicWorkspaceSettingRequest{Name: portFwdServerlessSetting}).
+				Return(tc.setting, tc.readErr)
+
+			assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, m.WorkspaceClient))
+		})
+	}
+}