direct: retry 504 errors

NEXT_CHANGELOG.md

@@ -8,5 +8,6 @@
 * `experimental open` now opens every DABs resource type that has a workspace URL, picking up `catalogs`, `schemas`, `volumes`, `database_instances`, `database_catalogs`, `synced_database_tables`, `postgres_catalogs`, `postgres_synced_tables`, `quality_monitors`, `vector_search_endpoints`, and `vector_search_indexes` ([#5346](https://github.com/databricks/cli/pull/5346)).
 
 ### Bundles
+* Retry transient HTTP 504 Gateway Timeout errors in direct deployment engine ([#5349](https://github.com/databricks/cli/pull/5349)).
 
 ### Dependency updates

acceptance/bin/fault.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+"""Set up a fault rule on the testserver for the current test token.
+
+Usage: fault.py PATTERN STATUS_CODE OFFSET TIMES
+
+  PATTERN     HTTP method and path, supports trailing * wildcard,
+              e.g. "PUT /api/2.0/permissions/pipelines/*"
+  STATUS_CODE HTTP status code to return, e.g. 504
+  OFFSET      number of requests to let through before fault starts
+  TIMES       number of times to return the fault response
+
+The rule is scoped to the current DATABRICKS_TOKEN so it only affects
+the test that registers it, even when tests share a server.
+"""
+
+import json
+import os
+import sys
+import urllib.request
+
+host = os.environ.get("DATABRICKS_HOST", "")
+token = os.environ.get("DATABRICKS_TOKEN", "")
+
+if not host:
+    print("DATABRICKS_HOST not set", file=sys.stderr)
+    sys.exit(1)
+
+if len(sys.argv) != 5:
+    print(f"usage: {sys.argv[0]} PATTERN STATUS_CODE OFFSET TIMES", file=sys.stderr)
+    sys.exit(1)
+
+pattern, status_code, offset, times = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), int(sys.argv[4])
+body = '{"error_code":"INJECTED","message":"Fault injected by test."}'
+
+data = json.dumps(
+    {
+        "pattern": pattern,
+        "status_code": status_code,
+        "body": body,
+        "offset": offset,
+        "times": times,
+    }
+).encode()
+
+req = urllib.request.Request(
+    f"{host}/__testserver/fault",
+    data=data,
+    headers={"Content-Type": "application/json", "Authorization": f"Bearer {token}"},
+    method="POST",
+)
+urllib.request.urlopen(req)

acceptance/bundle/resources/permissions/pipelines/504/create/databricks.yml

@@ -0,0 +1,10 @@
+bundle:
+  name: test-bundle
+
+resources:
+  pipelines:
+    foo:
+      name: foo
+      permissions:
+        - level: CAN_VIEW
+          user_name: viewer@example.com

acceptance/bundle/resources/permissions/pipelines/504/create/output.txt

@@ -0,0 +1,9 @@
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Deploying resources...
+Warn: deploying resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from PUT /api/2.0/permissions/pipelines/[UUID]
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py //api/2.0/permissions/pipelines
+"PUT /api/2.0/permissions/pipelines/[UUID]"
+"PUT /api/2.0/permissions/pipelines/[UUID]"

acceptance/bundle/resources/permissions/pipelines/504/create/script

@@ -0,0 +1,8 @@
+# Inject a single 504 on the first permissions PUT to simulate a transient error.
+# Permissions Set is idempotent, so DoCreate opts in via retrySafe and the deploy succeeds.
+fault.py "PUT /api/2.0/permissions/pipelines/*" 504 0 1
+
+$CLI bundle deploy
+
+# Two PUT requests should appear: the initial 504 and the successful retry.
+trace print_requests.py //api/2.0/permissions/pipelines | jq '.method + " " + .path'

acceptance/bundle/resources/permissions/pipelines/504/create/test.toml

@@ -0,0 +1 @@
+RecordRequests = true

acceptance/bundle/resources/permissions/pipelines/504/plan/databricks.yml

@@ -0,0 +1,10 @@
+bundle:
+  name: test-bundle
+
+resources:
+  pipelines:
+    foo:
+      name: foo
+      permissions:
+        - level: CAN_VIEW
+          user_name: viewer@example.com

acceptance/bundle/resources/permissions/pipelines/504/plan/output.txt

@@ -0,0 +1,31 @@
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] bundle plan
+Warn: planning resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from GET /api/2.0/permissions/pipelines/[UUID]
+Plan: 0 to add, 0 to change, 0 to delete, 2 unchanged
+
+>>> print_requests.py //api/2.0/permissions/pipelines --get --oneline
+2 {"method": "GET", "path": "/api/2.0/permissions/pipelines/[UUID]"}
+
+>>> [CLI] bundle plan
+Warn: planning resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from GET /api/2.0/permissions/pipelines/[UUID]
+Warn: planning resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from GET /api/2.0/permissions/pipelines/[UUID]
+Plan: 0 to add, 0 to change, 0 to delete, 2 unchanged
+3 {"method": "GET", "path": "/api/2.0/permissions/pipelines/[UUID]"}
+
+>>> musterr [CLI] bundle plan
+Warn: planning resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from GET /api/2.0/permissions/pipelines/[UUID]
+Warn: planning resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from GET /api/2.0/permissions/pipelines/[UUID]
+Error: cannot plan resources.pipelines.foo.permissions: reading id="/pipelines/[UUID]": Fault injected by test. (504 INJECTED)
+
+Endpoint: GET [DATABRICKS_URL]/api/2.0/permissions/pipelines/[UUID]?
+HTTP Status: 504 Gateway Timeout
+API error_code: INJECTED
+API message: Fault injected by test.
+
+Error: planning failed
+
+3 {"method": "GET", "path": "/api/2.0/permissions/pipelines/[UUID]"}

acceptance/bundle/resources/permissions/pipelines/504/plan/script

@@ -0,0 +1,15 @@
+# Deploy first so the permissions resource exists; plan reads it on subsequent runs.
+$CLI bundle deploy
+rm -f out.requests.txt
+
+fault.py "GET /api/2.0/permissions/pipelines/*" 504 0 1
+trace $CLI bundle plan
+trace print_requests.py //api/2.0/permissions/pipelines --get --oneline | uniq -c | sed 's/^ *//' | contains.py "2 "
+
+fault.py "GET /api/2.0/permissions/pipelines/*" 504 0 2
+trace $CLI bundle plan
+print_requests.py //api/2.0/permissions/pipelines --get --oneline | uniq -c | sed 's/^ *//' | contains.py "3 "
+
+fault.py "GET /api/2.0/permissions/pipelines/*" 504 0 3
+trace musterr $CLI bundle plan
+print_requests.py //api/2.0/permissions/pipelines --get --oneline | uniq -c | sed 's/^ *//' | contains.py "3 "

acceptance/bundle/resources/permissions/pipelines/504/test.toml

@@ -0,0 +1,4 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+EnvMatrix.READPLAN = []
+RecordRequests = true
+Env.DATABRICKS_BUNDLE_RETRY_INTERVAL_MS = "100"

acceptance/bundle/resources/permissions/pipelines/504/update/databricks.yml

@@ -0,0 +1,10 @@
+bundle:
+  name: test-bundle
+
+resources:
+  pipelines:
+    foo:
+      name: foo
+      permissions:
+        - level: CAN_VIEW
+          user_name: viewer@example.com

acceptance/bundle/resources/permissions/pipelines/504/update/output.txt

@@ -0,0 +1,17 @@
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Deploying resources...
+Warn: deploying resources.pipelines.foo.permissions: retrying after 504 Gateway Timeout from PUT /api/2.0/permissions/pipelines/[UUID]
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py //api/2.0/permissions/pipelines
+"PUT /api/2.0/permissions/pipelines/[UUID]"
+"PUT /api/2.0/permissions/pipelines/[UUID]"

acceptance/bundle/resources/permissions/pipelines/504/update/script

@@ -0,0 +1,13 @@
+trace $CLI bundle deploy
+
+update_file.py databricks.yml CAN_VIEW CAN_MANAGE
+
+# Inject a single 504 on the first permissions PUT to simulate a transient error.
+# The retrying adapter should retry after DATABRICKS_BUNDLE_RETRY_INTERVAL_MS and succeed.
+fault.py "PUT /api/2.0/permissions/pipelines/*" 504 0 1
+
+rm out.requests.txt
+trace $CLI bundle deploy
+
+# Two PUT requests should appear: the initial 504 and the successful retry.
+trace print_requests.py //api/2.0/permissions/pipelines | jq '.method + " " + .path'

bundle/direct/apply.go

@@ -8,12 +8,14 @@ import (
 	"reflect"
 
 	"github.com/databricks/cli/bundle/deployplan"
+	"github.com/databricks/cli/bundle/direct/dresources"
 	"github.com/databricks/cli/bundle/direct/dstate"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/apierr"
 )
 
 func (d *DeploymentUnit) Destroy(ctx context.Context, db *dstate.DeploymentState) error {
+	ctx = log.WithPrefix(ctx, "destroying "+d.ResourceKey)
 	id := db.GetResourceID(d.ResourceKey)
 	if id == "" {
 		log.Infof(ctx, "Cannot delete %s: missing from state", d.ResourceKey)
@@ -24,6 +26,7 @@ func (d *DeploymentUnit) Destroy(ctx context.Context, db *dstate.DeploymentState
 }
 
 func (d *DeploymentUnit) Deploy(ctx context.Context, db *dstate.DeploymentState, newState any, actionType deployplan.ActionType, planEntry *deployplan.PlanEntry) error {
+	ctx = log.WithPrefix(ctx, "deploying "+d.ResourceKey)
 	if actionType == deployplan.Create {
 		return d.Create(ctx, db, newState)
 	}
@@ -48,7 +51,18 @@ func (d *DeploymentUnit) Deploy(ctx context.Context, db *dstate.DeploymentState,
 }
 
 func (d *DeploymentUnit) Create(ctx context.Context, db *dstate.DeploymentState, newState any) error {
-	newID, remoteState, err := d.Adapter.DoCreate(ctx, newState)
+	var newID string
+	var remoteState any
+	_, err := retryWith(ctx, func(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)
+	}, func() (struct{}, error) {
+		var e error
+		newID, remoteState, e = d.Adapter.DoCreate(ctx, newState)
+		return struct{}{}, e
+	})
+	err = dresources.UnwrapRetrySafe(err)
 	if err != nil {
 		// No need to prefix error, there is no ambiguity (only one operation - DoCreate) and no additional context (like id)
 		return err
@@ -66,7 +80,9 @@ func (d *DeploymentUnit) Create(ctx context.Context, db *dstate.DeploymentState,
 		return fmt.Errorf("saving state after creating id=%s: %w", newID, err)
 	}
 
-	waitRemoteState, err := d.Adapter.WaitAfterCreate(ctx, newID, newState)
+	waitRemoteState, err := retryOnTransient(ctx, func() (any, error) {
+		return d.Adapter.WaitAfterCreate(ctx, newID, newState)
+	})
 	if err != nil {
 		return fmt.Errorf("waiting after creating id=%s: %w", newID, err)
 	}
@@ -89,7 +105,7 @@ 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 = d.Adapter.DoDelete(ctx, oldID, oldState)
+	err = retryOnTransientErr(ctx, func() error { return d.Adapter.DoDelete(ctx, oldID, oldState) })
 	if err != nil && !isResourceGone(err) && !isManagedByParent(err) {
 		return fmt.Errorf("deleting old id=%s: %w", oldID, err)
 	}
@@ -118,7 +134,9 @@ 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 := d.Adapter.DoUpdate(ctx, id, newState, planEntry)
+	remoteState, err := retryOnTransient(ctx, func() (any, error) {
+		return d.Adapter.DoUpdate(ctx, id, newState, planEntry)
+	})
 	if err != nil {
 		return fmt.Errorf("updating id=%s: %w", id, err)
 	}
@@ -133,7 +151,9 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState,
 		return fmt.Errorf("saving state id=%s: %w", id, err)
 	}
 
-	waitRemoteState, err := d.Adapter.WaitAfterUpdate(ctx, id, newState)
+	waitRemoteState, err := retryOnTransient(ctx, func() (any, error) {
+		return d.Adapter.WaitAfterUpdate(ctx, id, newState)
+	})
 	if err != nil {
 		return fmt.Errorf("waiting after updating id=%s: %w", id, err)
 	}
@@ -148,7 +168,13 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState,
 }
 
 func (d *DeploymentUnit) UpdateWithID(ctx context.Context, db *dstate.DeploymentState, oldID string, newState any) error {
-	newID, remoteState, err := d.Adapter.DoUpdateWithID(ctx, oldID, newState)
+	var newID string
+	var remoteState any
+	err := retryOnTransientErr(ctx, func() error {
+		var e error
+		newID, remoteState, e = d.Adapter.DoUpdateWithID(ctx, oldID, newState)
+		return e
+	})
 	if err != nil {
 		return fmt.Errorf("updating id=%s: %w", oldID, err)
 	}
@@ -169,7 +195,9 @@ func (d *DeploymentUnit) UpdateWithID(ctx context.Context, db *dstate.Deployment
 		return fmt.Errorf("saving state id=%s: %w", oldID, err)
 	}
 
-	waitRemoteState, err := d.Adapter.WaitAfterUpdate(ctx, newID, newState)
+	waitRemoteState, err := retryOnTransient(ctx, func() (any, error) {
+		return d.Adapter.WaitAfterUpdate(ctx, newID, newState)
+	})
 	if err != nil {
 		return fmt.Errorf("waiting after updating id=%s: %w", newID, err)
 	}
@@ -219,7 +247,7 @@ func (d *DeploymentUnit) Delete(ctx context.Context, db *dstate.DeploymentState,
 }
 
 func (d *DeploymentUnit) Resize(ctx context.Context, db *dstate.DeploymentState, id string, newState any) error {
-	err := d.Adapter.DoResize(ctx, id, newState)
+	err := retryOnTransientErr(ctx, func() error { return d.Adapter.DoResize(ctx, id, newState) })
 	if err != nil {
 		return fmt.Errorf("resizing id=%s: %w", id, err)
 	}
@@ -263,7 +291,9 @@ func (d *DeploymentUnit) refreshRemoteState(ctx context.Context, id string) erro
 	if d.RemoteState != nil {
 		return nil
 	}
-	remoteState, err := d.Adapter.DoRead(ctx, id)
+	remoteState, err := retryOnTransient(ctx, func() (any, error) {
+		return d.Adapter.DoRead(ctx, id)
+	})
 	if err != nil {
 		return fmt.Errorf("failed to refresh remote state id=%s: %w", id, err)
 	}

bundle/direct/bundle_plan.go

@@ -122,6 +122,7 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks
 	// We're processing resources in DAG order because we're resolving references (that can be resolved at plan stage).
 	g.Run(defaultParallelism, func(resourceKey string, failedDependency *string) bool {
 		errorPrefix := "cannot plan " + resourceKey
+		ctx := log.WithPrefix(ctx, "planning "+resourceKey)
 
 		entry, err := plan.WriteLockEntry(resourceKey)
 		if err != nil {
@@ -155,7 +156,9 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks
 				return false
 			}
 
-			remoteState, err := adapter.DoRead(ctx, id)
+			remoteState, err := retryOnTransient(ctx, func() (any, error) {
+				return adapter.DoRead(ctx, id)
+			})
 			if err != nil {
 				if isResourceGone(err) {
 					// no such resource
@@ -210,7 +213,9 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks
 			return false
 		}
 
-		remoteState, err := adapter.DoRead(ctx, dbentry.ID)
+		remoteState, err := retryOnTransient(ctx, func() (any, error) {
+			return adapter.DoRead(ctx, dbentry.ID)
+		})
 		if err != nil {
 			if isResourceGone(err) {
 				remoteState = nil

bundle/direct/dresources/grants.go

@@ -110,7 +110,8 @@ func (r *ResourceGrants) DoRead(ctx context.Context, id string) (*GrantsState, e
 func (r *ResourceGrants) DoCreate(ctx context.Context, state *GrantsState) (string, *GrantsState, error) {
 	_, err := r.DoUpdate(ctx, "", state, nil)
 	if err != nil {
-		return "", nil, err
+		// Grants Update is idempotent (additive PATCH), so retrying on transient errors is safe.
+		return "", nil, retrySafe(err)
 	}
 
 	return state.SecurableType + "/" + state.FullName, nil, nil

bundle/direct/dresources/permissions.go

@@ -219,7 +219,8 @@ func (r *ResourcePermissions) DoCreate(ctx context.Context, newState *Permission
 	// should we remember the default here?
 	_, err := r.DoUpdate(ctx, newState.ObjectID, newState, nil)
 	if err != nil {
-		return "", nil, err
+		// Permissions Set is idempotent (PUT), so retrying on transient errors is safe.
+		return "", nil, retrySafe(err)
 	}
 
 	return newState.ObjectID, nil, nil