diff --git a/.nextchanges/bundles/file-upload-retries.md b/.nextchanges/bundles/file-upload-retries.md new file mode 100644 index 00000000000..364761f4dd7 --- /dev/null +++ b/.nextchanges/bundles/file-upload-retries.md @@ -0,0 +1 @@ +* `databricks bundle deploy`, `databricks bundle sync` and `databricks sync` now retry file, library, state and metadata uploads that fail with a transient server error (5xx or 408), instead of failing. ([#7014](https://github.com/databricks/cli/pull/7014)) diff --git a/acceptance/bundle/deploy/files/upload-504/databricks.yml b/acceptance/bundle/deploy/files/upload-504/databricks.yml new file mode 100644 index 00000000000..576d7a9ef25 --- /dev/null +++ b/acceptance/bundle/deploy/files/upload-504/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: test-bundle diff --git a/acceptance/bundle/deploy/files/upload-504/out.test.toml b/acceptance/bundle/deploy/files/upload-504/out.test.toml new file mode 100644 index 00000000000..a927a5fbc06 --- /dev/null +++ b/acceptance/bundle/deploy/files/upload-504/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DMS = ["", "true"] diff --git a/acceptance/bundle/deploy/files/upload-504/output.txt b/acceptance/bundle/deploy/files/upload-504/output.txt new file mode 100644 index 00000000000..52ff524b127 --- /dev/null +++ b/acceptance/bundle/deploy/files/upload-504/output.txt @@ -0,0 +1,6 @@ +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Warn: retrying after 504 Gateway Timeout from POST /api/2.0/workspace/import: failed to upload /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml: Fault injected by test. +Files: 1 uploaded, 0 deleted +Resources: 0 created, 0 changed, 0 deleted, 0 unchanged +/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml +/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml diff --git a/acceptance/bundle/deploy/files/upload-504/script b/acceptance/bundle/deploy/files/upload-504/script new file mode 100644 index 00000000000..78740b9becc --- /dev/null +++ b/acceptance/bundle/deploy/files/upload-504/script @@ -0,0 +1,8 @@ +# Let the two deploy lock requests through, then inject a single 504 on the next import request, which is a file upload. +# Uploads overwrite, so they are retried and the deploy succeeds. +fault.py "POST /api/2.0/workspace/import" 504 2 1 + +$CLI bundle deploy + +# databricks.yml is the only file. Two requests for it should appear: the initial 504 and the successful retry. +print_requests.py //api/2.0/workspace/import | jq -r .body.multipart_form.path | grep '/files/' diff --git a/acceptance/bundle/deploy/files/upload-504/test.toml b/acceptance/bundle/deploy/files/upload-504/test.toml new file mode 100644 index 00000000000..f3ca35f85ab --- /dev/null +++ b/acceptance/bundle/deploy/files/upload-504/test.toml @@ -0,0 +1,2 @@ +RecordRequests = true +Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100" diff --git a/acceptance/bundle/sync/output.txt b/acceptance/bundle/sync/output.txt index 6ff42f41918..ca0b5b7258d 100644 --- a/acceptance/bundle/sync/output.txt +++ b/acceptance/bundle/sync/output.txt @@ -1,6 +1,6 @@ >>> [CLI] bundle sync --output text -Action: PUT: .gitignore, databricks.yml, dryrun/databricks.yml, dryrun/out.test.toml, ignored-folder/folder1/script.py, ignored-folder/script.py +Action: PUT: .gitignore, databricks.yml, dryrun/databricks.yml, dryrun/out.test.toml, ignored-folder/folder1/script.py, ignored-folder/script.py, upload-504/databricks.yml, upload-504/out.test.toml Initial Sync Complete Uploaded .gitignore Uploaded databricks.yml @@ -10,3 +10,6 @@ Uploaded dryrun/out.test.toml Uploaded ignored-folder/folder1 Uploaded ignored-folder/folder1/script.py Uploaded ignored-folder/script.py +Uploaded upload-504 +Uploaded upload-504/databricks.yml +Uploaded upload-504/out.test.toml diff --git a/acceptance/bundle/sync/upload-504/databricks.yml b/acceptance/bundle/sync/upload-504/databricks.yml new file mode 100644 index 00000000000..576d7a9ef25 --- /dev/null +++ b/acceptance/bundle/sync/upload-504/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: test-bundle diff --git a/acceptance/bundle/sync/upload-504/out.test.toml b/acceptance/bundle/sync/upload-504/out.test.toml new file mode 100644 index 00000000000..a927a5fbc06 --- /dev/null +++ b/acceptance/bundle/sync/upload-504/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DMS = ["", "true"] diff --git a/acceptance/bundle/sync/upload-504/output.txt b/acceptance/bundle/sync/upload-504/output.txt new file mode 100644 index 00000000000..af7934522d1 --- /dev/null +++ b/acceptance/bundle/sync/upload-504/output.txt @@ -0,0 +1,6 @@ +Action: PUT: databricks.yml +Warn: retrying after 504 Gateway Timeout from POST /api/2.0/workspace/import: failed to upload /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml: Fault injected by test. +Uploaded databricks.yml +Initial Sync Complete +/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml +/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/databricks.yml diff --git a/acceptance/bundle/sync/upload-504/script b/acceptance/bundle/sync/upload-504/script new file mode 100644 index 00000000000..d4e9e5d1b51 --- /dev/null +++ b/acceptance/bundle/sync/upload-504/script @@ -0,0 +1,7 @@ +# Inject a single 504 on the first file upload. Uploads overwrite, so they are retried and the sync succeeds. +fault.py "POST /api/2.0/workspace/import" 504 0 1 + +$CLI bundle sync + +# databricks.yml is the only file. Two requests for it should appear: the initial 504 and the successful retry. +print_requests.py //api/2.0/workspace/import | jq -r .body.multipart_form.path diff --git a/acceptance/bundle/sync/upload-504/test.toml b/acceptance/bundle/sync/upload-504/test.toml new file mode 100644 index 00000000000..f3ca35f85ab --- /dev/null +++ b/acceptance/bundle/sync/upload-504/test.toml @@ -0,0 +1,2 @@ +RecordRequests = true +Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100" diff --git a/acceptance/bundle/upload/internal_server_error/output.txt b/acceptance/bundle/upload/internal_server_error/output.txt index 794a0ed2072..2ae3634f94b 100644 --- a/acceptance/bundle/upload/internal_server_error/output.txt +++ b/acceptance/bundle/upload/internal_server_error/output.txt @@ -1,4 +1,6 @@ Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Warn: retrying after 500 Internal Server Error from POST /api/2.0/workspace/import: failed to upload /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/file_to_upload.txt: Fault injected by test. +Warn: retrying after 500 Internal Server Error from POST /api/2.0/workspace/import: failed to upload /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/file_to_upload.txt: Fault injected by test. Error: failed to upload /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/file_to_upload.txt: Fault injected by test. (500 INJECTED) Endpoint: POST [DATABRICKS_URL]/api/2.0/workspace/import diff --git a/acceptance/bundle/upload/internal_server_error/test.toml b/acceptance/bundle/upload/internal_server_error/test.toml new file mode 100644 index 00000000000..2889b56fb03 --- /dev/null +++ b/acceptance/bundle/upload/internal_server_error/test.toml @@ -0,0 +1 @@ +Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100" diff --git a/acceptance/cmd/sync/output.txt b/acceptance/cmd/sync/output.txt index a48c5ed331c..4695359a9d9 100644 --- a/acceptance/cmd/sync/output.txt +++ b/acceptance/cmd/sync/output.txt @@ -1,6 +1,6 @@ >>> [CLI] sync . /Users/[USERNAME] -Action: PUT: .gitignore, dryrun-missing-remote/hello.py, dryrun-missing-remote/out.test.toml, dryrun-missing-remote/output.txt, dryrun/out.test.toml, dryrun/output.txt, project-folder/app.py, project-folder/app.yaml, project-folder/query.sql +Action: PUT: .gitignore, dryrun-missing-remote/hello.py, dryrun-missing-remote/out.test.toml, dryrun-missing-remote/output.txt, dryrun/out.test.toml, dryrun/output.txt, project-folder/app.py, project-folder/app.yaml, project-folder/query.sql, upload-504/hello.py, upload-504/out.test.toml, upload-504/output.txt Initial Sync Complete Uploaded .gitignore Uploaded dryrun @@ -14,6 +14,10 @@ Uploaded project-folder Uploaded project-folder/app.py Uploaded project-folder/app.yaml Uploaded project-folder/query.sql +Uploaded upload-504 +Uploaded upload-504/hello.py +Uploaded upload-504/out.test.toml +Uploaded upload-504/output.txt >>> [CLI] sync . /Users/[USERNAME] --exclude project-folder/app.* Action: DELETE: project-folder/app.py, project-folder/app.yaml diff --git a/acceptance/cmd/sync/upload-504/hello.py b/acceptance/cmd/sync/upload-504/hello.py new file mode 100644 index 00000000000..11b15b1a458 --- /dev/null +++ b/acceptance/cmd/sync/upload-504/hello.py @@ -0,0 +1 @@ +print("hello") diff --git a/acceptance/cmd/sync/upload-504/out.test.toml b/acceptance/cmd/sync/upload-504/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/sync/upload-504/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/sync/upload-504/output.txt b/acceptance/cmd/sync/upload-504/output.txt new file mode 100644 index 00000000000..cb9b61d2b0c --- /dev/null +++ b/acceptance/cmd/sync/upload-504/output.txt @@ -0,0 +1,6 @@ +Action: PUT: hello.py +Warn: retrying after 504 Gateway Timeout from POST /api/2.0/workspace/import: failed to upload /Users/[USERNAME]/sync-504-[UNIQUE_NAME]/hello.py: Fault injected by test. +Uploaded hello.py +Initial Sync Complete +/Users/[USERNAME]/sync-504-[UNIQUE_NAME]/hello.py +/Users/[USERNAME]/sync-504-[UNIQUE_NAME]/hello.py diff --git a/acceptance/cmd/sync/upload-504/script b/acceptance/cmd/sync/upload-504/script new file mode 100644 index 00000000000..3be71acff67 --- /dev/null +++ b/acceptance/cmd/sync/upload-504/script @@ -0,0 +1,24 @@ +# .gitignore excludes itself so that hello.py is the only uploaded file and the injected fault always hits it. +mkdir .git +cat > .gitignore << EOF2 +.databricks/ +.gitignore +EOF2 + +# Use an ephemeral remote directory so concurrent test runs do not interfere. +remote_dir="/Users/$CURRENT_USER_NAME/sync-504-$UNIQUE_NAME" + +cleanup() { + MSYS_NO_PATHCONV=1 $CLI workspace delete --recursive "$remote_dir" 2>/dev/null || true + rm -rf .git .gitignore .databricks "$OUT_REQUESTS" +} +trap cleanup EXIT + +# Inject a single 504 on the first file upload. Uploads overwrite, so they are retried and the sync succeeds. +fault.py "POST /api/2.0/workspace/import" 504 0 1 + +# MSYS_NO_PATHCONV=1 prevents Git Bash on Windows from converting /Users/... to C:/Program Files/Git/Users/... +MSYS_NO_PATHCONV=1 $CLI sync . "$remote_dir" + +# Two requests for hello.py should appear: the initial 504 and the successful retry. +print_requests.py //api/2.0/workspace/import | jq -r .body.multipart_form.path diff --git a/acceptance/cmd/sync/upload-504/test.toml b/acceptance/cmd/sync/upload-504/test.toml new file mode 100644 index 00000000000..f3ca35f85ab --- /dev/null +++ b/acceptance/cmd/sync/upload-504/test.toml @@ -0,0 +1,2 @@ +RecordRequests = true +Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100" diff --git a/bundle/deploy/filer.go b/bundle/deploy/filer.go index 8fdbda90e01..8d7b2a57636 100644 --- a/bundle/deploy/filer.go +++ b/bundle/deploy/filer.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/retry" "github.com/databricks/databricks-sdk-go" ) @@ -72,7 +73,7 @@ func StateFiler(ctx context.Context, b *bundle.Bundle) (filer.Filer, error) { } return stateFiler{ - filer: f, + filer: retry.Filer(f), root: filer.NewWorkspaceRootPath(b.Config.Workspace.StatePath), workspaceClient: w, }, nil diff --git a/bundle/deploy/metadata/upload.go b/bundle/deploy/metadata/upload.go index 79546518a36..3217b054034 100644 --- a/bundle/deploy/metadata/upload.go +++ b/bundle/deploy/metadata/upload.go @@ -9,6 +9,7 @@ import ( "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/retry" ) const metadataFileName = "metadata.json" @@ -28,10 +29,11 @@ func (m *upload) Name() string { } func (m *upload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { - f, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), b.Config.Workspace.StatePath) + workspaceFiler, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), b.Config.Workspace.StatePath) if err != nil { return diag.FromErr(err) } + f := retry.Filer(workspaceFiler) metadata, err := json.MarshalIndent(b.Metadata, "", " ") if err != nil { diff --git a/bundle/direct/apply.go b/bundle/direct/apply.go index 41ffd4bcb8c..57367883c6c 100644 --- a/bundle/direct/apply.go +++ b/bundle/direct/apply.go @@ -12,6 +12,7 @@ import ( "github.com/databricks/cli/bundle/direct/dresources" "github.com/databricks/cli/bundle/direct/dstate" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/retry" "github.com/databricks/databricks-sdk-go/apierr" ) @@ -62,10 +63,10 @@ func (d *DeploymentUnit) Deploy(ctx context.Context, db *dstate.DeploymentState, func (d *DeploymentUnit) Create(ctx context.Context, db *dstate.DeploymentState, newState any) error { var newID string var remoteState any - _, err := retryWith(ctx, func(err error) bool { + _, err := retry.With(ctx, func(ctx context.Context, err error) bool { // For DoCreate, retry feature is opt-in via retrySafe(err) error wrapper _, ok := errors.AsType[*dresources.RetrySafeError](err) - return ok && isTransient(ctx, err) + return ok && retry.IsTransient(ctx, err) }, func() (struct{}, error) { var e error newID, remoteState, e = d.Adapter.DoCreate(ctx, newState) @@ -90,7 +91,7 @@ func (d *DeploymentUnit) Create(ctx context.Context, db *dstate.DeploymentState, } waitRemoteState, err := waitCapped(ctx, d.MaxWait, "creation of "+d.ResourceKey, func(ctx context.Context) (any, error) { - return retryOnTransient(ctx, func() (any, error) { + return retry.With(ctx, retry.IsTransient, func() (any, error) { return d.Adapter.WaitAfterCreate(ctx, newID, newState) }) }) @@ -116,7 +117,9 @@ func (d *DeploymentUnit) Recreate(ctx context.Context, db *dstate.DeploymentStat // MANAGED_BY_PARENT is still disregarded — the subsequent Create with // replace_existing=true will reconfigure the parent-managed resource in // place, matching the Terraform provider's recreate behaviour. - err = retryOnTransientErr(ctx, func() error { return d.Adapter.DoDelete(ctx, oldID, oldState) }) + _, err = retry.With(ctx, retry.IsTransient, func() (struct{}, error) { + return struct{}{}, d.Adapter.DoDelete(ctx, oldID, oldState) + }) if err != nil && !apierr.IsMissing(err) && !isManagedByParent(err) { if !d.deleteConfirmedGone(ctx, oldID) { return fmt.Errorf("deleting old id=%s: %w", oldID, err) @@ -154,7 +157,7 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState, return fmt.Errorf("internal error: DoUpdate not implemented for resource %s", d.ResourceKey) } - remoteState, err := retryOnTransient(ctx, func() (any, error) { + remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return d.Adapter.DoUpdate(ctx, id, newState, planEntry) }) if err != nil { @@ -186,7 +189,7 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState, } } - waitRemoteState, err := retryOnTransient(ctx, func() (any, error) { + waitRemoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return d.Adapter.WaitAfterUpdate(ctx, id, newState) }) if err != nil { @@ -205,10 +208,10 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState, func (d *DeploymentUnit) UpdateWithID(ctx context.Context, db *dstate.DeploymentState, oldID string, newState any, planEntry *deployplan.PlanEntry) error { var newID string var remoteState any - err := retryOnTransientErr(ctx, func() error { + _, err := retry.With(ctx, retry.IsTransient, func() (struct{}, error) { var e error newID, remoteState, e = d.Adapter.DoUpdateWithID(ctx, oldID, newState, planEntry) - return e + return struct{}{}, e }) if err != nil { return fmt.Errorf("updating id=%s: %w", oldID, err) @@ -230,7 +233,7 @@ func (d *DeploymentUnit) UpdateWithID(ctx context.Context, db *dstate.Deployment return fmt.Errorf("saving state id=%s: %w", oldID, err) } - waitRemoteState, err := retryOnTransient(ctx, func() (any, error) { + waitRemoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return d.Adapter.WaitAfterUpdate(ctx, newID, newState) }) if err != nil { @@ -311,7 +314,9 @@ func (d *DeploymentUnit) deleteConfirmedGone(ctx context.Context, id string) boo } func (d *DeploymentUnit) Resize(ctx context.Context, db *dstate.DeploymentState, id string, newState any, entry *deployplan.PlanEntry) error { - err := retryOnTransientErr(ctx, func() error { return d.Adapter.DoResize(ctx, id, newState, entry) }) + _, err := retry.With(ctx, retry.IsTransient, func() (struct{}, error) { + return struct{}{}, d.Adapter.DoResize(ctx, id, newState, entry) + }) if err != nil { return fmt.Errorf("resizing id=%s: %w", id, err) } @@ -370,7 +375,7 @@ func (d *DeploymentUnit) refreshRemoteState(ctx context.Context, id string) erro if d.RemoteState != nil { return nil } - remoteState, err := retryOnTransient(ctx, func() (any, error) { + remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return d.Adapter.DoRead(ctx, id) }) if err != nil { diff --git a/bundle/direct/bundle_plan.go b/bundle/direct/bundle_plan.go index d403ca53f76..98da4f1376d 100644 --- a/bundle/direct/bundle_plan.go +++ b/bundle/direct/bundle_plan.go @@ -19,6 +19,7 @@ import ( "github.com/databricks/cli/libs/dyn/dynvar" "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/logdiag" + "github.com/databricks/cli/libs/retry" "github.com/databricks/cli/libs/structs/structaccess" "github.com/databricks/cli/libs/structs/structdiff" "github.com/databricks/cli/libs/structs/structpath" @@ -236,7 +237,7 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks return true } - remoteState, err := retryOnTransient(ctx, func() (any, error) { + remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return adapter.DoRead(ctx, id) }) if err != nil { @@ -314,7 +315,7 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks return false } - remoteState, err := retryOnTransient(ctx, func() (any, error) { + remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) { return adapter.DoRead(ctx, dbentry.ID) }) if err != nil { diff --git a/bundle/direct/maxwait.go b/bundle/direct/maxwait.go index a13a5faf61d..6a9fff26ac9 100644 --- a/bundle/direct/maxwait.go +++ b/bundle/direct/maxwait.go @@ -19,7 +19,7 @@ import ( const maxWaitUnset = time.Duration(-1) // resourceMaxWait returns the cap on waiting for a resource to reach its target state, or -// maxWaitUnset when the environment variable is absent. Unlike retryInterval, a malformed +// maxWaitUnset when the environment variable is absent. Unlike the retry interval, a malformed // value is an error rather than a silent fallback: ignoring a typo would restore the // multi-hour default wait that the user was trying to shorten. func resourceMaxWait(ctx context.Context) (time.Duration, error) { @@ -86,7 +86,7 @@ func waitCapped[T any](ctx context.Context, maxWait time.Duration, description s // isWaitTimeout reports whether err is a wait that ran out of time rather than a resource // that failed. Two shapes reach here: retries.Poll reports a deadline as ErrTimedOut wrapping -// the last poll message, while retryWith returns a bare context error when the deadline lands +// the last poll message, while retry.With returns a bare context error when the deadline lands // while it sleeps between transient-error retries. func isWaitTimeout(err error) bool { if _, ok := errors.AsType[*retries.ErrTimedOut](err); ok { diff --git a/bundle/direct/maxwait_test.go b/bundle/direct/maxwait_test.go index 4b3f8e2b546..da4fcdeaa43 100644 --- a/bundle/direct/maxwait_test.go +++ b/bundle/direct/maxwait_test.go @@ -64,7 +64,7 @@ func TestWaitCappedAbandonsPoll(t *testing.T) { assert.Positive(t, polls) } -// TestWaitCappedAbandonsBareContextError covers the other shape: retryWith returns ctx.Err() +// TestWaitCappedAbandonsBareContextError covers the other shape: retry.With returns ctx.Err() // directly when the cap lands while it sleeps between transient-error retries. func TestWaitCappedAbandonsBareContextError(t *testing.T) { _, err := waitCapped(t.Context(), 20*time.Millisecond, "test resource", waitForCtx) diff --git a/bundle/direct/retry.go b/bundle/direct/retry.go deleted file mode 100644 index 43f5d6f0a4b..00000000000 --- a/bundle/direct/retry.go +++ /dev/null @@ -1,85 +0,0 @@ -package direct - -import ( - "context" - "errors" - "fmt" - "net/http" - "strconv" - "time" - - bundleenv "github.com/databricks/cli/bundle/env" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go/apierr" -) - -const maxRetries = 2 - -var defaultRetryInterval = 15 * time.Second - -func retryInterval(ctx context.Context) time.Duration { - v, ok := bundleenv.RetryIntervalMs(ctx) - if !ok { - return defaultRetryInterval - } - ms, err := strconv.Atoi(v) - if err != nil { - return defaultRetryInterval - } - return time.Duration(ms) * time.Millisecond -} - -// isTransient returns true for 504 errors that the SDK did not already retry. -// The SDK retries 504s matching its allTransientErrors patterns (e.g. "deadline exceeded"); -// this covers the remaining 504s like TEMPORARILY_UNAVAILABLE. -func isTransient(ctx context.Context, err error) bool { - apiErr, ok := errors.AsType[*apierr.APIError](err) - if !ok { - return false - } - if apiErr.IsRetriable(ctx) { - // Already handled by SDK - return false - } - return apiErr.StatusCode >= 500 || apiErr.StatusCode == http.StatusRequestTimeout -} - -// retryWith retries fn while check returns true for the error, up to maxRetries times. -func retryWith[T any](ctx context.Context, check func(error) bool, fn func() (T, error)) (T, error) { - interval := retryInterval(ctx) - for attempt := 0; ; attempt++ { - result, err := fn() - if err == nil || attempt >= maxRetries || !check(err) { - return result, err - } - msg := "retrying" - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { - endpoint := "" - if rw := apiErr.ResponseWrapper; rw != nil && rw.Response != nil && rw.Response.Request != nil { - req := rw.Response.Request - endpoint = fmt.Sprintf(" from %s %s", req.Method, req.URL.Path) - } - msg = fmt.Sprintf("retrying after %d %s%s", apiErr.StatusCode, http.StatusText(apiErr.StatusCode), endpoint) - } - log.Warnf(ctx, "%s", msg) - select { - case <-ctx.Done(): - var zero T - return zero, ctx.Err() - case <-time.After(interval): - } - } -} - -// retryOnTransient retries fn on transient 504 errors that the SDK did not already handle. -func retryOnTransient[T any](ctx context.Context, fn func() (T, error)) (T, error) { - return retryWith(ctx, func(err error) bool { return isTransient(ctx, err) }, fn) -} - -// retryOnTransientErr wraps retryOnTransient for functions that return only an error. -func retryOnTransientErr(ctx context.Context, fn func() error) error { - _, err := retryOnTransient(ctx, func() (struct{}, error) { - return struct{}{}, fn() - }) - return err -} diff --git a/bundle/env/retry_interval.go b/bundle/env/retry_interval.go deleted file mode 100644 index d00d88e0b25..00000000000 --- a/bundle/env/retry_interval.go +++ /dev/null @@ -1,11 +0,0 @@ -package env - -import "context" - -// RetryIntervalMsVariable names the environment variable that overrides the retry interval for bundle operations. -const RetryIntervalMsVariable = "DATABRICKS_BUNDLE_RETRY_INTERVAL_MS" - -// RetryIntervalMs returns the retry interval override (in milliseconds) for bundle operations. -func RetryIntervalMs(ctx context.Context) (string, bool) { - return get(ctx, []string{RetryIntervalMsVariable}) -} diff --git a/bundle/libraries/upload.go b/bundle/libraries/upload.go index 247c380a864..eeb71cfb9e8 100644 --- a/bundle/libraries/upload.go +++ b/bundle/libraries/upload.go @@ -15,6 +15,7 @@ import ( "github.com/databricks/cli/libs/dyn" "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/retry" "golang.org/x/sync/errgroup" ) @@ -69,7 +70,7 @@ func (u *upload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { } cmdio.LogProgress(ctx, fmt.Sprintf("Uploading %s...", relPath)) errs.Go(func() error { - return UploadFile(errCtx, source, u.client) + return UploadFile(errCtx, source, retry.Filer(u.client)) }) } diff --git a/libs/retry/filer.go b/libs/retry/filer.go new file mode 100644 index 00000000000..091b4c73952 --- /dev/null +++ b/libs/retry/filer.go @@ -0,0 +1,57 @@ +package retry + +import ( + "bytes" + "context" + "io" + "slices" + + "github.com/databricks/cli/libs/filer" +) + +// retryingFiler retries overwriting writes that fail with a transient server error. +type retryingFiler struct { + filer.Filer +} + +// Filer wraps f so that Write retries on transient server errors. +// Only writes with [filer.OverwriteIfExists] are retried: they are idempotent, whereas repeating +// a write that does not overwrite could fail on the file its own earlier attempt created. +// All other methods are passed through. +func Filer(f filer.Filer) filer.Filer { + return &retryingFiler{Filer: f} +} + +func (r *retryingFiler) Write(ctx context.Context, path string, reader io.Reader, mode ...filer.WriteMode) error { + if !slices.Contains(mode, filer.OverwriteIfExists) { + return r.Filer.Write(ctx, path, reader, mode...) + } + + // Every attempt needs the full contents, so rewind the reader or buffer it. + rewind := func() (io.Reader, error) { return reader, nil } + if seeker, ok := reader.(io.ReadSeeker); ok { + start, err := seeker.Seek(0, io.SeekCurrent) + if err != nil { + return err + } + rewind = func() (io.Reader, error) { + _, err := seeker.Seek(start, io.SeekStart) + return reader, err + } + } else { + body, err := io.ReadAll(reader) + if err != nil { + return err + } + rewind = func() (io.Reader, error) { return bytes.NewReader(body), nil } + } + + _, err := With(ctx, IsTransient, func() (struct{}, error) { + attempt, err := rewind() + if err != nil { + return struct{}{}, err + } + return struct{}{}, r.Filer.Write(ctx, path, attempt, mode...) + }) + return err +} diff --git a/libs/retry/filer_test.go b/libs/retry/filer_test.go new file mode 100644 index 00000000000..a264ce47066 --- /dev/null +++ b/libs/retry/filer_test.go @@ -0,0 +1,83 @@ +package retry + +import ( + "bytes" + "context" + "io" + "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" +) + +type flakyFiler struct { + filer.Filer + failures int + err error + bodies []string +} + +func (f *flakyFiler) Write(ctx context.Context, path string, r io.Reader, mode ...filer.WriteMode) error { + b, err := io.ReadAll(r) + if err != nil { + return err + } + f.bodies = append(f.bodies, string(b)) + if len(f.bodies) <= f.failures { + return f.err + } + return nil +} + +var errTransient = &apierr.APIError{StatusCode: http.StatusGatewayTimeout, Message: "TEMPORARILY_UNAVAILABLE"} + +func TestFilerRetriesSeekableReader(t *testing.T) { + f := &flakyFiler{failures: 1, err: errTransient} + err := Filer(f).Write(testContext(t), "a", bytes.NewReader([]byte("hello")), filer.OverwriteIfExists) + require.NoError(t, err) + assert.Equal(t, []string{"hello", "hello"}, f.bodies) +} + +func TestFilerRetriesNonSeekableReader(t *testing.T) { + f := &flakyFiler{failures: 1, err: errTransient} + // io.NopCloser hides the Seek method. + r := io.NopCloser(strings.NewReader("hello")) + err := Filer(f).Write(testContext(t), "a", r, filer.OverwriteIfExists) + require.NoError(t, err) + assert.Equal(t, []string{"hello", "hello"}, f.bodies) +} + +func TestFilerRewindsToInitialReaderPosition(t *testing.T) { + f := &flakyFiler{failures: 1, err: errTransient} + r := strings.NewReader("xxhello") + _, err := r.Seek(2, io.SeekStart) + require.NoError(t, err) + err = Filer(f).Write(testContext(t), "a", r, filer.OverwriteIfExists) + require.NoError(t, err) + assert.Equal(t, []string{"hello", "hello"}, f.bodies) +} + +func TestFilerGivesUp(t *testing.T) { + f := &flakyFiler{failures: 100, err: errTransient} + err := Filer(f).Write(testContext(t), "a", strings.NewReader("x"), filer.OverwriteIfExists) + require.Error(t, err) + assert.Len(t, f.bodies, MaxRetries+1) +} + +func TestFilerDoesNotRetryOtherErrors(t *testing.T) { + f := &flakyFiler{failures: 100, err: &apierr.APIError{StatusCode: http.StatusForbidden}} + err := Filer(f).Write(testContext(t), "a", strings.NewReader("x"), filer.OverwriteIfExists) + require.Error(t, err) + assert.Len(t, f.bodies, 1) +} + +func TestFilerDoesNotRetryWithoutOverwrite(t *testing.T) { + f := &flakyFiler{failures: 100, err: errTransient} + err := Filer(f).Write(testContext(t), "a", strings.NewReader("x")) + require.Error(t, err) + assert.Len(t, f.bodies, 1) +} diff --git a/libs/retry/retry.go b/libs/retry/retry.go new file mode 100644 index 00000000000..81b894cbb88 --- /dev/null +++ b/libs/retry/retry.go @@ -0,0 +1,87 @@ +// Package retry retries calls that fail with transient server errors the SDK does not retry itself. +package retry + +import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/apierr" +) + +// MaxRetries is the number of additional attempts made after the first one fails. +const MaxRetries = 2 + +// IntervalMsVariable names the environment variable that overrides the pause between attempts. +const IntervalMsVariable = "DATABRICKS_BUNDLE_RETRY_INTERVAL_MS" + +const defaultInterval = 15 * time.Second + +// interval returns the pause between attempts: the override in milliseconds if set and valid, 15s otherwise. +func interval(ctx context.Context) time.Duration { + v, ok := env.Lookup(ctx, IntervalMsVariable) + if !ok { + return defaultInterval + } + ms, err := strconv.Atoi(v) + if err != nil { + return defaultInterval + } + return time.Duration(ms) * time.Millisecond +} + +// IsTransient returns true for 5xx and 408 errors that the SDK did not already retry. +// The SDK retries 504s matching its allTransientErrors patterns (e.g. "deadline exceeded"); +// this covers the remaining 504s like TEMPORARILY_UNAVAILABLE. +func IsTransient(ctx context.Context, err error) bool { + apiErr, ok := errors.AsType[*apierr.APIError](err) + if !ok { + return false + } + if apiErr.IsRetriable(ctx) { + // Already handled by SDK + return false + } + return apiErr.StatusCode >= 500 || apiErr.StatusCode == http.StatusRequestTimeout +} + +// With retries fn while check returns true for the error, up to MaxRetries times. +func With[T any](ctx context.Context, check func(context.Context, error) bool, fn func() (T, error)) (T, error) { + wait := interval(ctx) + for attempt := 0; ; attempt++ { + result, err := fn() + if err == nil || attempt >= MaxRetries || !check(ctx, err) { + return result, err + } + log.Warnf(ctx, "%s", retryMessage(err)) + select { + case <-ctx.Done(): + var zero T + return zero, ctx.Err() + case <-time.After(wait): + } + } +} + +// retryMessage describes the error that triggered a retry. +func retryMessage(err error) string { + msg := "retrying" + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + endpoint := "" + if rw := apiErr.ResponseWrapper; rw != nil && rw.Response != nil && rw.Response.Request != nil { + req := rw.Response.Request + endpoint = fmt.Sprintf(" from %s %s", req.Method, req.URL.Path) + } + msg = fmt.Sprintf("retrying after %d %s%s", apiErr.StatusCode, http.StatusText(apiErr.StatusCode), endpoint) + // Include the context the caller added to the error, e.g. the path of the file being written. + if err.Error() != apiErr.Error() { + msg += ": " + err.Error() + } + } + return msg +} diff --git a/libs/retry/retry_test.go b/libs/retry/retry_test.go new file mode 100644 index 00000000000..231389d4aa2 --- /dev/null +++ b/libs/retry/retry_test.go @@ -0,0 +1,95 @@ +package retry + +import ( + "context" + "errors" + "fmt" + "net/http" + "testing" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIsTransient(t *testing.T) { + ctx := t.Context() + assert.True(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusGatewayTimeout, Message: "TEMPORARILY_UNAVAILABLE"})) + assert.True(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusInternalServerError})) + assert.True(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusRequestTimeout})) + // Retried by the SDK already. + assert.False(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusServiceUnavailable})) + assert.False(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusGatewayTimeout, Message: "deadline exceeded"})) + assert.False(t, IsTransient(ctx, &apierr.APIError{StatusCode: http.StatusForbidden})) + assert.False(t, IsTransient(ctx, errors.New("boom"))) +} + +// testContext shortens the pause between attempts. +func testContext(t *testing.T) context.Context { + return env.Set(t.Context(), IntervalMsVariable, "1") +} + +func TestInterval(t *testing.T) { + ctx := t.Context() + assert.Equal(t, 15*time.Second, interval(ctx)) + assert.Equal(t, 100*time.Millisecond, interval(env.Set(ctx, IntervalMsVariable, "100"))) + assert.Equal(t, 15*time.Second, interval(env.Set(ctx, IntervalMsVariable, "abc"))) +} + +func TestWithRetriesUntilSuccess(t *testing.T) { + ctx := testContext(t) + calls := 0 + got, err := With(ctx, func(context.Context, error) bool { return true }, func() (int, error) { + calls++ + if calls <= MaxRetries { + return 0, errors.New("fail") + } + return 42, nil + }) + require.NoError(t, err) + assert.Equal(t, 42, got) + assert.Equal(t, MaxRetries+1, calls) +} + +func TestWithGivesUp(t *testing.T) { + ctx := testContext(t) + calls := 0 + _, err := With(ctx, func(context.Context, error) bool { return true }, func() (int, error) { + calls++ + return 0, errors.New("fail") + }) + require.EqualError(t, err, "fail") + assert.Equal(t, MaxRetries+1, calls) +} + +func TestWithDoesNotRetryWhenCheckFails(t *testing.T) { + ctx := testContext(t) + calls := 0 + _, err := With(ctx, func(context.Context, error) bool { return false }, func() (int, error) { + calls++ + return 0, errors.New("fail") + }) + require.Error(t, err) + assert.Equal(t, 1, calls) +} + +func TestWithStopsOnContextCancel(t *testing.T) { + ctx, cancel := context.WithCancel(env.Set(t.Context(), IntervalMsVariable, "3600000")) + calls := 0 + _, err := With(ctx, func(context.Context, error) bool { return true }, func() (int, error) { + calls++ + cancel() + return 0, errors.New("fail") + }) + require.ErrorIs(t, err, context.Canceled) + assert.Equal(t, 1, calls) +} + +func TestRetryMessage(t *testing.T) { + apiErr := &apierr.APIError{StatusCode: http.StatusGatewayTimeout, Message: "boom"} + assert.Equal(t, "retrying", retryMessage(errors.New("x"))) + assert.Equal(t, "retrying after 504 Gateway Timeout", retryMessage(apiErr)) + assert.Equal(t, "retrying after 504 Gateway Timeout: failed to upload /a/b: boom", retryMessage(fmt.Errorf("failed to upload /a/b: %w", apiErr))) +} diff --git a/libs/sync/sync.go b/libs/sync/sync.go index a78ff920eb5..1e31b85481f 100644 --- a/libs/sync/sync.go +++ b/libs/sync/sync.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/fileset" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/retry" "github.com/databricks/cli/libs/vfs" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/iam" @@ -114,10 +115,11 @@ func New(ctx context.Context, opts SyncOptions) (*Sync, error) { } } - filer, err := filer.NewWorkspaceFilesClient(opts.WorkspaceClient, opts.RemotePath) + workspaceFiler, err := filer.NewWorkspaceFilesClient(opts.WorkspaceClient, opts.RemotePath) if err != nil { return nil, err } + filer := retry.Filer(workspaceFiler) var notifier EventNotifier outputWaitGroup := &stdsync.WaitGroup{}