diff --git a/.nextchanges/cli/files-api-retry.md b/.nextchanges/cli/files-api-retry.md new file mode 100644 index 00000000000..7105a5bd296 --- /dev/null +++ b/.nextchanges/cli/files-api-retry.md @@ -0,0 +1 @@ +* Fixed `fs` commands on UC Volumes failing immediately on transient Files API errors (such as `temporarily unavailable`) instead of retrying. ([#7252](https://github.com/databricks/cli/pull/7252)) diff --git a/libs/filer/files_client.go b/libs/filer/files_client.go index 1638acb73bd..12819ce5ac0 100644 --- a/libs/filer/files_client.go +++ b/libs/filer/files_client.go @@ -16,7 +16,10 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/sdk-go/core/apierr" + "github.com/databricks/sdk-go/core/apiretry" + "github.com/databricks/sdk-go/core/ops" files "github.com/databricks/sdk-go/files/v2" + "github.com/databricks/sdk-go/options/call" "github.com/databricks/sdk-go/options/client" "golang.org/x/sync/errgroup" ) @@ -25,6 +28,32 @@ import ( // on a large-file part transfer. It matches the files/v2 engine's own default. const cloudResponseHeaderTimeout = 60 * time.Second +// files/v2 calls are not retried unless a retrier is passed per call, so without +// these options a transient 5xx (e.g. "GCP is temporarily unavailable") fails the +// command outright. The status codes match the files/v2 upload engine's +// control-plane retries; the timeout matches its control-plane retry budget and +// the default SDK retry timeout. Only idempotent calls get them: uploads cannot +// replay their body, and a retried delete could turn a success into a 404. +var ( + filesAPIRetryStatusCodes = []int{ + http.StatusRequestTimeout, + http.StatusTooManyRequests, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout, + } + filesAPIRetryTimeout = 5 * time.Minute +) + +func filesAPIRetryOptions() []call.Option { + return []call.Option{ + call.WithRetrier(func() ops.Retrier { + return apiretry.NewRetrier(ops.BackoffPolicy{}, apiretry.RetrierConfig{StatusCodes: filesAPIRetryStatusCodes}) + }), + call.WithTimeout(filesAPIRetryTimeout), + } +} + // httpStatus returns the HTTP status code of err if it is (or wraps) an // [apierr.APIError], and -1 otherwise. The filer keys its error mapping off the // HTTP status rather than the SDK's canonical codes.Code: the Files API reports @@ -213,7 +242,7 @@ func (w *FilesClient) Write(ctx context.Context, name string, reader io.Reader, // Check that target path exists if CreateParentDirectories mode is not set if !slices.Contains(mode, CreateParentDirectories) { dir := path.Dir(absPath) - _, err := w.client.GetDirectoryMetadata(ctx, files.GetDirectoryMetadataRequest{DirectoryPath: &dir}) + _, err := w.client.GetDirectoryMetadata(ctx, files.GetDirectoryMetadataRequest{DirectoryPath: &dir}, filesAPIRetryOptions()...) if err != nil { // This API returns a 404 if the directory doesn't exist. if httpStatus(err) == http.StatusNotFound { @@ -295,7 +324,7 @@ func (w *FilesClient) Read(ctx context.Context, name string) (io.ReadCloser, err return nil, err } - resp, err := w.client.DownloadFile(ctx, files.DownloadFileRequest{FilePath: &absPath}) + resp, err := w.client.DownloadFile(ctx, files.DownloadFileRequest{FilePath: &absPath}, filesAPIRetryOptions()...) // Return early on success. if err == nil { @@ -477,7 +506,7 @@ func (w *FilesClient) ReadDir(ctx context.Context, name string) ([]fs.DirEntry, var entries []fs.DirEntry for entry, err := range w.client.ListDirectoryContentsIter(ctx, files.ListDirectoryContentsRequest{ DirectoryPath: &absPath, - }) { + }, filesAPIRetryOptions()...) { if err != nil { // This API returns a 404 if the specified path does not exist. if httpStatus(err) == http.StatusNotFound { @@ -513,7 +542,7 @@ func (w *FilesClient) Mkdir(ctx context.Context, name string) error { return err } - _, err = w.client.CreateDirectory(ctx, files.CreateDirectoryRequest{DirectoryPath: &absPath}) + _, err = w.client.CreateDirectory(ctx, files.CreateDirectoryRequest{DirectoryPath: &absPath}, filesAPIRetryOptions()...) // This API returns a 409 when a file already exists at the path (the create // is not idempotent over a file). @@ -531,7 +560,7 @@ func (w *FilesClient) statFile(ctx context.Context, name string) (fs.FileInfo, e return nil, err } - resp, err := w.client.GetFileMetadata(ctx, files.GetFileMetadataRequest{FilePath: &absPath}) + resp, err := w.client.GetFileMetadata(ctx, files.GetFileMetadataRequest{FilePath: &absPath}, filesAPIRetryOptions()...) // If the HEAD requests succeeds, the file exists. if err == nil { @@ -557,7 +586,7 @@ func (w *FilesClient) statDir(ctx context.Context, name string) (fs.FileInfo, er return nil, err } - _, err = w.client.GetDirectoryMetadata(ctx, files.GetDirectoryMetadataRequest{DirectoryPath: &absPath}) + _, err = w.client.GetDirectoryMetadata(ctx, files.GetDirectoryMetadataRequest{DirectoryPath: &absPath}, filesAPIRetryOptions()...) // If the HEAD requests succeeds, the directory exists. if err == nil { diff --git a/libs/filer/files_client_test.go b/libs/filer/files_client_test.go index 3579c5464c5..27510a17bae 100644 --- a/libs/filer/files_client_test.go +++ b/libs/filer/files_client_test.go @@ -6,8 +6,10 @@ import ( "fmt" "io" "io/fs" + "net/http" "os" "path/filepath" + "sync/atomic" "testing" cliauth "github.com/databricks/cli/libs/auth" @@ -145,6 +147,49 @@ func TestFilesClientMkdirWhenFileExists(t *testing.T) { assert.ErrorIs(t, err, fs.ErrExist) } +// statDirWithFirstStatus fails the first directory HEAD with the given status +// and lets later ones succeed. It returns the number of HEAD requests the server +// saw and the Stat error. +func statDirWithFirstStatus(t *testing.T, statusCode int) (int, error) { + t.Helper() + + server := testserver.New(t) + var calls atomic.Int32 + server.Handle("HEAD", "/api/2.0/fs/directories/{path...}", func(req testserver.Request) any { + if calls.Add(1) == 1 { + return testserver.Response{ + StatusCode: statusCode, + Body: map[string]any{"error_code": "TEMPORARILY_UNAVAILABLE", "message": "GCP is temporarily unavailable; try again later."}, + } + } + return "" + }) + + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + + f, err := NewFilesClient(t.Context(), client, "/Volumes/main/schema/vol") + require.NoError(t, err) + + _, err = f.(*FilesClient).statDir(t.Context(), "dir") + return int(calls.Load()), err +} + +func TestFilesClientRetriesTransientError(t *testing.T) { + calls, err := statDirWithFirstStatus(t, http.StatusServiceUnavailable) + require.NoError(t, err) + assert.Equal(t, 2, calls) +} + +func TestFilesClientDoesNotRetryInternalError(t *testing.T) { + calls, err := statDirWithFirstStatus(t, http.StatusInternalServerError) + require.Error(t, err) + assert.Equal(t, 1, calls) +} + // onlyReader hides the Seek method of an underlying reader, modelling a // non-seekable stream (e.g. a remote download body). type onlyReader struct{ io.Reader }