Integrate deployment metadata service for locking and state

acceptance/bundle/deploy/metadata-service/databricks.yml

@@ -0,0 +1,7 @@
+bundle:
+  name: metadata-service-test
+
+resources:
+  jobs:
+    test_job:
+      name: test-job

acceptance/bundle/deploy/metadata-service/deploy-error/databricks.yml

@@ -0,0 +1,7 @@
+bundle:
+  name: metadata-service-error-test
+
+resources:
+  jobs:
+    test_job:
+      name: test-job

acceptance/bundle/deploy/metadata-service/deploy-error/output.txt

@@ -0,0 +1,61 @@
+
+>>> musterr [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/metadata-service-error-test/default/files...
+Deploying resources...
+Error: cannot create resources.jobs.test_job: Invalid job configuration. (400 INVALID_PARAMETER_VALUE)
+
+Endpoint: POST [DATABRICKS_URL]/api/2.2/jobs/create
+HTTP Status: 400 Bad Request
+API error_code: INVALID_PARAMETER_VALUE
+API message: Invalid job configuration.
+
+Updating deployment state...
+
+>>> print_requests.py --get //bundle
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments",
+  "q": {
+    "deployment_id": "[UUID]"
+  },
+  "body": {
+    "target_name": "default"
+  }
+}
+{
+  "method": "GET",
+  "path": "/api/2.0/bundle/deployments/[UUID]"
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions",
+  "q": {
+    "version_id": "1"
+  },
+  "body": {
+    "cli_version": "[DEV_VERSION]",
+    "version_type": "VERSION_TYPE_DEPLOY",
+    "target_name": "default"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/operations",
+  "q": {
+    "resource_key": "resources.jobs.test_job"
+  },
+  "body": {
+    "resource_key": "resources.jobs.test_job",
+    "action_type": "OPERATION_ACTION_TYPE_CREATE",
+    "status": "OPERATION_STATUS_FAILED",
+    "error_message": "Invalid job configuration."
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/complete",
+  "body": {
+    "name": "deployments/[UUID]/versions/1",
+    "completion_reason": "VERSION_COMPLETE_FAILURE"
+  }
+}

acceptance/bundle/deploy/metadata-service/deploy-error/script

@@ -0,0 +1,5 @@
+# Deploy with the metadata service enabled, expecting a resource creation failure.
+trace musterr $CLI bundle deploy
+
+# Print the metadata service requests to verify the failed operation is reported.
+trace print_requests.py --get //bundle

acceptance/bundle/deploy/metadata-service/deploy-error/test.toml

@@ -0,0 +1,8 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+EnvMatrix.DATABRICKS_BUNDLE_MANAGED_STATE = ["true"]
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /api/2.2/jobs/create"
+Response.StatusCode = 400
+Response.Body = '{"error_code": "INVALID_PARAMETER_VALUE", "message": "Invalid job configuration."}'

acceptance/bundle/deploy/metadata-service/output.txt

@@ -0,0 +1,117 @@
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/metadata-service-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py --get //bundle
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments",
+  "q": {
+    "deployment_id": "[UUID]"
+  },
+  "body": {
+    "target_name": "default"
+  }
+}
+{
+  "method": "GET",
+  "path": "/api/2.0/bundle/deployments/[UUID]"
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions",
+  "q": {
+    "version_id": "1"
+  },
+  "body": {
+    "cli_version": "[DEV_VERSION]",
+    "version_type": "VERSION_TYPE_DEPLOY",
+    "target_name": "default"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/operations",
+  "q": {
+    "resource_key": "resources.jobs.test_job"
+  },
+  "body": {
+    "resource_key": "resources.jobs.test_job",
+    "action_type": "OPERATION_ACTION_TYPE_CREATE",
+    "resource_id": "[NUMID]",
+    "status": "OPERATION_STATUS_SUCCEEDED"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/complete",
+  "body": {
+    "name": "deployments/[UUID]/versions/1",
+    "completion_reason": "VERSION_COMPLETE_SUCCESS"
+  }
+}
+
+>>> [CLI] bundle destroy --auto-approve
+The following resources will be deleted:
+  delete resources.jobs.test_job
+
+All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/metadata-service-test/default
+
+Deleting files...
+Destroy complete!
+
+>>> print_requests.py --get //bundle
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments",
+  "q": {
+    "deployment_id": "[UUID]"
+  },
+  "body": {
+    "target_name": "default"
+  }
+}
+{
+  "method": "GET",
+  "path": "/api/2.0/bundle/deployments/[UUID]"
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions",
+  "q": {
+    "version_id": "1"
+  },
+  "body": {
+    "cli_version": "[DEV_VERSION]",
+    "version_type": "VERSION_TYPE_DESTROY",
+    "target_name": "default"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/operations",
+  "q": {
+    "resource_key": "resources.jobs.test_job"
+  },
+  "body": {
+    "resource_key": "resources.jobs.test_job",
+    "action_type": "OPERATION_ACTION_TYPE_DELETE",
+    "resource_id": "[NUMID]",
+    "status": "OPERATION_STATUS_SUCCEEDED"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/bundle/deployments/[UUID]/versions/1/complete",
+  "body": {
+    "name": "deployments/[UUID]/versions/1",
+    "completion_reason": "VERSION_COMPLETE_SUCCESS"
+  }
+}
+{
+  "method": "DELETE",
+  "path": "/api/2.0/bundle/deployments/[UUID]"
+}

acceptance/bundle/deploy/metadata-service/script

@@ -0,0 +1,11 @@
+# Deploy with the metadata service enabled.
+trace $CLI bundle deploy
+
+# Print all metadata service requests made during deploy.
+trace print_requests.py --get //bundle
+
+# Destroy with the metadata service enabled.
+trace $CLI bundle destroy --auto-approve
+
+# Print all metadata service requests made during destroy.
+trace print_requests.py --get //bundle

acceptance/bundle/deploy/metadata-service/test.toml

@@ -0,0 +1,3 @@
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+EnvMatrix.DATABRICKS_BUNDLE_MANAGED_STATE = ["true"]
+RecordRequests = true

bundle/bundle.go

@@ -225,10 +225,10 @@ func TryLoad(ctx context.Context) *Bundle {
 	return b
 }
 
-func (b *Bundle) WorkspaceClientE() (*databricks.WorkspaceClient, error) {
+func (b *Bundle) WorkspaceClientE(ctx context.Context) (*databricks.WorkspaceClient, error) {
 	b.clientOnce.Do(func() {
 		var err error
-		b.client, err = b.Config.Workspace.Client()
+		b.client, err = b.Config.Workspace.Client(ctx)
 		if err != nil {
 			b.clientErr = fmt.Errorf("cannot resolve bundle auth configuration: %w", err)
 		}
@@ -238,7 +238,7 @@ func (b *Bundle) WorkspaceClientE() (*databricks.WorkspaceClient, error) {
 }
 
 func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
-	client, err := b.WorkspaceClientE()
+	client, err := b.WorkspaceClientE(context.TODO())
 	if err != nil {
 		panic(err)
 	}

bundle/bundle_test.go

@@ -184,20 +184,21 @@ func TestClearWorkspaceClient(t *testing.T) {
 	b.Config.Workspace.Host = "https://nonexistent.example.com"
 	b.Config.Workspace.Profile = "profile-A"
 
-	_, err1 := b.WorkspaceClientE()
+	ctx := t.Context()
+	_, err1 := b.WorkspaceClientE(ctx)
 	require.Error(t, err1)
 	assert.Contains(t, err1.Error(), "profile-A")
 
 	// Without retry, second call returns the same cached error (same object).
-	_, err1b := b.WorkspaceClientE()
+	_, err1b := b.WorkspaceClientE(ctx)
 	assert.Same(t, err1, err1b, "expected same cached error without retry")
 
 	// After retry, change the profile to "profile-B" and call again.
 	// If retry didn't re-execute, the error would still mention "profile-A".
 	b.ClearWorkspaceClient()
 	b.Config.Workspace.Profile = "profile-B"
 
-	_, err2 := b.WorkspaceClientE()
+	_, err2 := b.WorkspaceClientE(ctx)
 	require.Error(t, err2)
 	assert.Contains(t, err2.Error(), "profile-B", "expected re-execution to pick up new profile")
 	assert.NotContains(t, err2.Error(), "profile-A", "stale cached error should not appear")

bundle/config/workspace.go

@@ -1,11 +1,14 @@
 package config
 
 import (
+	"context"
+	"net/http"
 	"os"
 	"path/filepath"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/databrickscfg"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/marshal"
@@ -156,7 +159,7 @@ func (w *Workspace) NormalizeHostURL() {
 	}
 }
 
-func (w *Workspace) Client() (*databricks.WorkspaceClient, error) {
+func (w *Workspace) Client(ctx context.Context) (*databricks.WorkspaceClient, error) {
 	// Extract query parameters (?o=, ?a=) from the host URL before building
 	// the SDK config. This ensures workspace_id and account_id are available
 	// for profile resolution during EnsureResolved().
@@ -193,9 +196,34 @@ func (w *Workspace) Client() (*databricks.WorkspaceClient, error) {
 		}
 	}
 
+	// If DATABRICKS_LITESWAP_ID is set, wrap the transport to inject the
+	// x-databricks-traffic-id header for routing to the liteswap instance.
+	if liteswapID := env.Get(ctx, "DATABRICKS_LITESWAP_ID"); liteswapID != "" {
+		inner := cfg.HTTPTransport
+		if inner == nil {
+			inner = http.DefaultTransport
+		}
+		cfg.HTTPTransport = &liteswapTransport{
+			inner:     inner,
+			trafficID: "testenv://liteswap/" + liteswapID,
+		}
+	}
+
 	return databricks.NewWorkspaceClient((*databricks.Config)(cfg))
 }
 
+// liteswapTransport injects the x-databricks-traffic-id header to route
+// requests to a liteswap service instance.
+type liteswapTransport struct {
+	inner     http.RoundTripper
+	trafficID string
+}
+
+func (t *liteswapTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set("x-databricks-traffic-id", t.trafficID)
+	return t.inner.RoundTrip(req)
+}
+
 func init() {
 	arg0 := os.Args[0]