experimental/ssh: retry tunnel binary upload on transient errors

NEXT_CHANGELOG.md

@@ -7,6 +7,7 @@
 ### CLI
 
 * An explicitly selected profile (`--profile` or a bundle's `workspace.profile`) now takes precedence over auth environment variables (`DATABRICKS_HOST`, `DATABRICKS_TOKEN`, etc.) instead of being silently shadowed by them; env vars still fill auth fields the profile leaves empty ([#5096](https://github.com/databricks/cli/issues/5096)).
+* `databricks ssh connect` now retries the SSH tunnel binary upload on transient workspace-files errors (a cold `get-status` request timeout or a 5xx), instead of aborting the whole connect on a single blip ([#5836](https://github.com/databricks/cli/pull/5836)).
 
 ### Bundles
 

experimental/ssh/internal/client/releases.go

@@ -10,14 +10,21 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/databricks/cli/experimental/ssh/internal/workspace"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/retries"
 	"golang.org/x/net/http2"
 )
 
+// uploadRetryTimeout bounds how long the transient-error retries for a single
+// workspace-files operation (stat or upload) may run before giving up.
+const uploadRetryTimeout = 5 * time.Minute
+
 type releaseProvider func(ctx context.Context, architecture, version, releasesDir string) (io.ReadCloser, error)
 
 func UploadTunnelReleases(ctx context.Context, client *databricks.WorkspaceClient, version, releasesDir string) error {
@@ -47,39 +54,95 @@ func uploadReleases(ctx context.Context, workspaceFiler filer.Filer, getRelease
 		remoteBinaryPath := filepath.ToSlash(filepath.Join(remoteSubFolder, "databricks"))
 		remoteArchivePath := filepath.ToSlash(filepath.Join(remoteSubFolder, "databricks.zip"))
 
-		_, err := workspaceFiler.Stat(ctx, remoteBinaryPath)
-		if err == nil {
+		exists, err := binaryExists(ctx, workspaceFiler, remoteBinaryPath)
+		if err != nil {
+			return fmt.Errorf("failed to check if file %s exists in workspace: %w", remoteBinaryPath, err)
+		}
+		if exists {
 			log.Infof(ctx, "File %s already exists in the workspace, skipping upload", remoteBinaryPath)
 			continue
-		} else if !errors.Is(err, fs.ErrNotExist) {
-			return fmt.Errorf("failed to check if file %s exists in workspace: %w", remoteBinaryPath, err)
 		}
 
+		log.Infof(ctx, "Uploading %s to the workspace", fileName)
+		if err := uploadRelease(ctx, workspaceFiler, getRelease, arch, version, releasesDir, remoteArchivePath); err != nil {
+			return err
+		}
+		log.Infof(ctx, "Successfully uploaded %s to workspace", remoteBinaryPath)
+	}
+
+	return nil
+}
+
+// binaryExists reports whether the tunnel binary is already present in the workspace,
+// retrying the stat on transient errors. The workspace-files get-status endpoint can
+// stall (a cold request routinely takes ~60s, right at the SDK's per-request timeout)
+// or return a transient 5xx, and a single such failure otherwise aborts the whole connect.
+func binaryExists(ctx context.Context, workspaceFiler filer.Filer, remoteBinaryPath string) (bool, error) {
+	var exists bool
+	err := retries.Wait(ctx, uploadRetryTimeout, func() *retries.Err {
+		_, statErr := workspaceFiler.Stat(ctx, remoteBinaryPath)
+		switch {
+		case statErr == nil:
+			exists = true
+			return nil
+		case errors.Is(statErr, fs.ErrNotExist):
+			exists = false
+			return nil
+		case isRetriableUploadError(ctx, statErr):
+			return retries.Continue(statErr)
+		default:
+			return retries.Halt(statErr)
+		}
+	})
+	return exists, err
+}
+
+// uploadRelease uploads a single architecture's release archive, retrying transient
+// failures. The reader is re-fetched on each attempt because Write drains it.
+func uploadRelease(ctx context.Context, workspaceFiler filer.Filer, getRelease releaseProvider, arch, version, releasesDir, remoteArchivePath string) error {
+	return retries.Wait(ctx, uploadRetryTimeout, func() *retries.Err {
 		releaseReader, err := getRelease(ctx, arch, version, releasesDir)
 		if err != nil {
-			return fmt.Errorf("failed to get archive for architecture %s: %w", arch, err)
+			return retries.Halt(fmt.Errorf("failed to get archive for architecture %s: %w", arch, err))
 		}
 		defer releaseReader.Close()
 
-		log.Infof(ctx, "Uploading %s to the workspace", fileName)
 		// workspace-files/import-file API will automatically unzip the payload,
 		// producing the filerRoot/remoteSubFolder/*archive-contents* structure, with 'databricks' binary inside.
 		err = workspaceFiler.Write(ctx, remoteArchivePath, releaseReader, filer.OverwriteIfExists, filer.CreateParentDirectories)
-		if err != nil {
-			if isStreamResetError(err) {
-				return fmt.Errorf("failed to upload file %s to workspace: %w\n\n"+
-					"The connection was closed before the upload finished. "+
-					"This is usually caused by a network intermediary (corporate egress proxy, VPN, or firewall/WAF) "+
-					"enforcing a request-body size limit on POSTs to *.cloud.databricks.com. "+
-					"Try running this command from a network without such restrictions",
-					remoteArchivePath, err)
-			}
-			return fmt.Errorf("failed to upload file %s to workspace: %w", remoteArchivePath, err)
+		switch {
+		case err == nil:
+			return nil
+		case isStreamResetError(err):
+			// A stream reset is a proxy body-size rejection, not a transient blip: retrying
+			// won't help, so fail with the actionable hint instead of exhausting the budget.
+			return retries.Halt(fmt.Errorf("failed to upload file %s to workspace: %w\n\n"+
+				"The connection was closed before the upload finished. "+
+				"This is usually caused by a network intermediary (corporate egress proxy, VPN, or firewall/WAF) "+
+				"enforcing a request-body size limit on POSTs to *.cloud.databricks.com. "+
+				"Try running this command from a network without such restrictions",
+				remoteArchivePath, err))
+		case isRetriableUploadError(ctx, err):
+			return retries.Continue(fmt.Errorf("failed to upload file %s to workspace: %w", remoteArchivePath, err))
+		default:
+			return retries.Halt(fmt.Errorf("failed to upload file %s to workspace: %w", remoteArchivePath, err))
 		}
-		log.Infof(ctx, "Successfully uploaded %s to workspace", remoteBinaryPath)
-	}
+	})
+}
 
-	return nil
+// isRetriableUploadError reports whether a workspace-files stat/upload error is worth
+// retrying: a transient API error (per the SDK's own classification, e.g. 429/503) or a
+// timeout/deadline, which the workspace-files get-status endpoint produces on a cold request.
+func isRetriableUploadError(ctx context.Context, err error) bool {
+	if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && apiErr.IsRetriable(ctx) {
+		return true
+	}
+	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, os.ErrDeadlineExceeded) {
+		return true
+	}
+	// The SDK surfaces a per-request inactivity timeout as a plain error whose message
+	// ends with "request timed out after ...", losing any typed deadline value.
+	return strings.Contains(err.Error(), "request timed out after")
 }
 
 // isStreamResetError reports whether err looks like an HTTP/2 stream reset from

experimental/ssh/internal/client/releases_test.go

@@ -1,14 +1,54 @@
 package client
 
 import (
+	"context"
 	"errors"
 	"fmt"
+	"io"
+	"io/fs"
+	"net/http"
+	"strings"
 	"testing"
 
+	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/net/http2"
 )
 
+// stubFiler is a filer.Filer whose Stat and Write return scripted results per call,
+// so tests can drive the retry loops in binaryExists/uploadRelease deterministically.
+type stubFiler struct {
+	filer.Filer
+	statErrs  []error
+	writeErrs []error
+	statCalls int
+	writes    int
+}
+
+func (s *stubFiler) Stat(ctx context.Context, name string) (fs.FileInfo, error) {
+	err := s.statErrs[s.statCalls]
+	s.statCalls++
+	return nil, err
+}
+
+func (s *stubFiler) Write(ctx context.Context, path string, reader io.Reader, mode ...filer.WriteMode) error {
+	// Drain the reader as the real filer would, so a retry must supply a fresh one.
+	_, _ = io.Copy(io.Discard, reader)
+	err := s.writeErrs[s.writes]
+	s.writes++
+	return err
+}
+
+func fakeRelease(ctx context.Context, arch, version, releasesDir string) (io.ReadCloser, error) {
+	return io.NopCloser(strings.NewReader("archive")), nil
+}
+
+func timeoutErr() error {
+	return errors.New(`Get "https://example.test/api/2.0/workspace/get-status": request timed out after 1m0s of inactivity`)
+}
+
 func TestIsStreamResetError(t *testing.T) {
 	tests := []struct {
 		name string
@@ -43,3 +83,76 @@
 		})
 	}
 }
+
+func TestIsRetriableUploadError(t *testing.T) {
+	ctx := context.Background()
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"inactivity timeout string", timeoutErr(), true},
+		{"context deadline", context.DeadlineExceeded, true},
+		{"retriable API 503", &apierr.APIError{StatusCode: http.StatusServiceUnavailable}, true},
+		{"retriable API 429", &apierr.APIError{StatusCode: http.StatusTooManyRequests}, true},
+		{"non-retriable API 404", &apierr.APIError{StatusCode: http.StatusNotFound}, false},
+		{"plain error", errors.New("connection refused"), false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, isRetriableUploadError(ctx, tt.err))
+		})
+	}
+}
+
+func TestBinaryExistsRetriesTransientStat(t *testing.T) {
+	ctx := context.Background()
+
+	// A transient timeout on the first stat, then a clean "exists" on the retry.
+	f := &stubFiler{statErrs: []error{timeoutErr(), nil}}
+	exists, err := binaryExists(ctx, f, "amd64/databricks")
+	require.NoError(t, err)
+	assert.True(t, exists)
+	assert.Equal(t, 2, f.statCalls)
+}
+
+func TestBinaryExistsNotFoundNoRetry(t *testing.T) {
+	ctx := context.Background()
+
+	// A definitive "not found" must resolve immediately (it's the signal to upload), not retry.
+	f := &stubFiler{statErrs: []error{fs.ErrNotExist}}
+	exists, err := binaryExists(ctx, f, "amd64/databricks")
+	require.NoError(t, err)
+	assert.False(t, exists)
+	assert.Equal(t, 1, f.statCalls)
+}
+
+func TestBinaryExistsHaltsOnNonRetriable(t *testing.T) {
+	ctx := context.Background()
+
+	f := &stubFiler{statErrs: []error{&apierr.APIError{StatusCode: http.StatusForbidden, Message: "denied"}}}
+	_, err := binaryExists(ctx, f, "amd64/databricks")
+	require.Error(t, err)
+	assert.Equal(t, 1, f.statCalls)
+}
+
+func TestUploadReleaseRetriesTransientWrite(t *testing.T) {
+	ctx := context.Background()
+
+	// First write fails transiently; the retry must re-fetch a fresh reader and succeed.
+	f := &stubFiler{writeErrs: []error{&apierr.APIError{StatusCode: http.StatusServiceUnavailable}, nil}}
+	err := uploadRelease(ctx, f, fakeRelease, "amd64", "1.0.0", "", "amd64/databricks.zip")
+	require.NoError(t, err)
+	assert.Equal(t, 2, f.writes)
+}
+
+func TestUploadReleaseStreamResetNoRetry(t *testing.T) {
+	ctx := context.Background()
+
+	// A stream reset is a proxy body-size rejection: fail fast with the hint, don't retry.
+	f := &stubFiler{writeErrs: []error{http2.StreamError{StreamID: 1, Code: http2.ErrCodeNo}}}
+	err := uploadRelease(ctx, f, fakeRelease, "amd64", "1.0.0", "", "amd64/databricks.zip")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "request-body size limit")
+	assert.Equal(t, 1, f.writes)
+}