Retry transient Files API errors in fs commands on UC Volumes

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

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 {

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 }