From 33cd676e1df909ac756d8b1ca535867f00c7088e Mon Sep 17 00:00:00 2001 From: TanishqDatabricks Date: Mon, 6 Jul 2026 10:57:58 +0200 Subject: [PATCH] experimental/ssh: retry tunnel binary upload on transient errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The binary-upload phase of ssh connect did a single Stat and single Write against the workspace-files API, so one transient failure aborted the whole connect. In practice this fires often: a cold get-status request routinely takes ~60s (right at the SDK's per-request inactivity timeout), and the import endpoint intermittently returns 5xx. Wrap the stat and the upload in a bounded retry (5 min) that re-runs on transient errors — SDK-classified retriable API errors (429/503/...) and request timeouts. The reader is re-fetched each attempt since Write drains it. fs.ErrNotExist (the signal to upload) and stream resets (a proxy body-size rejection, where retrying won't help) are not retried. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + experimental/ssh/internal/client/releases.go | 101 +++++++++++++--- .../ssh/internal/client/releases_test.go | 113 ++++++++++++++++++ 3 files changed, 196 insertions(+), 19 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 2c07236abf0..8bd45164701 100644 --- a/NEXT_CHANGELOG.md +++ b/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 diff --git a/experimental/ssh/internal/client/releases.go b/experimental/ssh/internal/client/releases.go index 58a8e5796a1..a890aced4b6 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/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 diff --git a/experimental/ssh/internal/client/releases_test.go b/experimental/ssh/internal/client/releases_test.go index 6444a779c60..9b998e501c3 100644 --- a/experimental/ssh/internal/client/releases_test.go +++ b/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 TestIsStreamResetError(t *testing.T) { }) } } + +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) +}