Retry bundle file uploads on transient server errors

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

acceptance/bundle/deploy/files/upload-504/databricks.yml

@@ -0,0 +1,2 @@
+bundle:
+  name: test-bundle

acceptance/bundle/deploy/files/upload-504/out.test.toml

@@ -0,0 +1,2 @@
+Cloud = false
+EnvMatrix.DMS = ["", "true"]

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

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/'

acceptance/bundle/deploy/files/upload-504/test.toml

@@ -0,0 +1,2 @@
+RecordRequests = true
+Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100"

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

acceptance/bundle/sync/upload-504/databricks.yml

@@ -0,0 +1,2 @@
+bundle:
+  name: test-bundle

acceptance/bundle/sync/upload-504/out.test.toml

@@ -0,0 +1,2 @@
+Cloud = false
+EnvMatrix.DMS = ["", "true"]

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

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

acceptance/bundle/sync/upload-504/test.toml

@@ -0,0 +1,2 @@
+RecordRequests = true
+Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100"

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

acceptance/bundle/upload/internal_server_error/test.toml

@@ -0,0 +1 @@
+Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100"

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

acceptance/cmd/sync/upload-504/hello.py

@@ -0,0 +1 @@
+print("hello")

acceptance/cmd/sync/upload-504/out.test.toml

@@ -0,0 +1 @@
+Cloud = false

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

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

acceptance/cmd/sync/upload-504/test.toml

@@ -0,0 +1,2 @@
+RecordRequests = true
+Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100"

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

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 {

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 {

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 {

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 {

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)