direct: Handle num_workers and workload_type fields for clusters as in TF (#3834)

andrewnester · web-flow · commit 14e8e4376b69 · 2025-10-29T14:07:01.000+01:00
## Changes
Handle num_workers and workload_type fields for clusters as in TF

## Why
To avoid differences in behaviour between 2 implementations
1. When `autoscale` is not set, `num_workers` should be sent even if
it's zero
2. when `workload_type` is set, both `notebook` and `job` clients
default to `true` unless overridden

## Tests
Added an acceptance test
diff --git a/acceptance/bundle/resources/clusters/deploy/num_workers_absent/databricks.yml b/acceptance/bundle/resources/clusters/deploy/num_workers_absent/databricks.yml
@@ -0,0 +1,23 @@
+bundle:
+  name: test-deploy-cluster-num-workers-absent
+
+resources:
+  clusters:
+    # Expecting num_workers to be sent as zero
+    cluster1:
+      cluster_name: test-cluster-1
+      spark_version: 13.3.x-scala2.12
+      node_type_id: i3.xlarge
+      spark_conf:
+        "spark.executor.memory": "2g"
+
+    # Expecting num_workers to be sent as absent
+    cluster2:
+      cluster_name: test-cluster-2
+      spark_version: 13.3.x-scala2.12
+      node_type_id: i3.xlarge
+      autoscale:
+        min_workers: 2
+        max_workers: 7
+      spark_conf:
+        "spark.executor.memory": "2g"
diff --git a/acceptance/bundle/resources/clusters/deploy/num_workers_absent/out.test.toml b/acceptance/bundle/resources/clusters/deploy/num_workers_absent/out.test.toml
diff --git a/acceptance/bundle/resources/clusters/deploy/num_workers_absent/output.txt b/acceptance/bundle/resources/clusters/deploy/num_workers_absent/output.txt
@@ -0,0 +1,39 @@
+
+>>> errcode [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-deploy-cluster-num-workers-absent/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py //clusters/create
+{
+  "method": "POST",
+  "path": "/api/2.1/clusters/create",
+  "body": {
+    "autotermination_minutes": 60,
+    "cluster_name": "test-cluster-1",
+    "node_type_id": "[NODE_TYPE_ID]",
+    "num_workers": 0,
+    "spark_conf": {
+      "spark.executor.memory": "2g"
+    },
+    "spark_version": "13.3.x-scala2.12"
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.1/clusters/create",
+  "body": {
+    "autoscale": {
+      "max_workers": 7,
+      "min_workers": 2
+    },
+    "autotermination_minutes": 60,
+    "cluster_name": "test-cluster-2",
+    "node_type_id": "[NODE_TYPE_ID]",
+    "spark_conf": {
+      "spark.executor.memory": "2g"
+    },
+    "spark_version": "13.3.x-scala2.12"
+  }
+}
diff --git a/acceptance/bundle/resources/clusters/deploy/num_workers_absent/script b/acceptance/bundle/resources/clusters/deploy/num_workers_absent/script
@@ -0,0 +1,2 @@
+trace errcode $CLI bundle deploy
+trace print_requests.py //clusters/create | jq -s 'sort_by(.body.cluster_name)[]'
diff --git a/acceptance/bundle/resources/clusters/deploy/num_workers_absent/test.toml b/acceptance/bundle/resources/clusters/deploy/num_workers_absent/test.toml
@@ -0,0 +1,11 @@
+Local = true
+Cloud = false
+RecordRequests = true
+
+Ignore = [
+    "databricks.yml",
+]
+
+[[Repls]]
+Old = "[0-9]{4}-[0-9]{6}-[0-9a-z]{8}"
+New = "[CLUSTER-ID]"
diff --git a/acceptance/bundle/resources/clusters/deploy/workload_type/databricks.yml b/acceptance/bundle/resources/clusters/deploy/workload_type/databricks.yml
@@ -0,0 +1,15 @@
+bundle:
+  name: test-deploy-cluster-workload-type
+
+resources:
+  clusters:
+    cluster1:
+      cluster_name: test-cluster
+      spark_version: 13.3.x-scala2.12
+      node_type_id: i3.xlarge
+      num_workers: 2
+      workload_type:
+        clients:
+          jobs: false
+      spark_conf:
+        "spark.executor.memory": "2g"
diff --git a/acceptance/bundle/resources/clusters/deploy/workload_type/out.test.toml b/acceptance/bundle/resources/clusters/deploy/workload_type/out.test.toml
diff --git a/acceptance/bundle/resources/clusters/deploy/workload_type/output.txt b/acceptance/bundle/resources/clusters/deploy/workload_type/output.txt
@@ -0,0 +1,28 @@
+
+>>> errcode [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-deploy-cluster-workload-type/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+>>> print_requests.py //clusters/create
+{
+  "method": "POST",
+  "path": "/api/2.1/clusters/create",
+  "body": {
+    "autotermination_minutes": 60,
+    "cluster_name": "test-cluster",
+    "node_type_id": "[NODE_TYPE_ID]",
+    "num_workers": 2,
+    "spark_conf": {
+      "spark.executor.memory": "2g"
+    },
+    "spark_version": "13.3.x-scala2.12",
+    "workload_type": {
+      "clients": {
+        "jobs": false,
+        "notebooks": true
+      }
+    }
+  }
+}
diff --git a/acceptance/bundle/resources/clusters/deploy/workload_type/script b/acceptance/bundle/resources/clusters/deploy/workload_type/script
@@ -0,0 +1,2 @@
+trace errcode $CLI bundle deploy
+trace print_requests.py //clusters/create
diff --git a/acceptance/bundle/resources/clusters/deploy/workload_type/test.toml b/acceptance/bundle/resources/clusters/deploy/workload_type/test.toml
@@ -0,0 +1,11 @@
+Local = true
+Cloud = false
+RecordRequests = true
+
+Ignore = [
+    "databricks.yml",
+]
+
+[[Repls]]
+Old = "[0-9]{4}-[0-9]{6}-[0-9a-z]{8}"
+New = "[CLUSTER-ID]"
diff --git a/bundle/config/mutator/resourcemutator/resource_mutator.go b/bundle/config/mutator/resourcemutator/resource_mutator.go
@@ -96,6 +96,8 @@ func applyInitializeMutators(ctx context.Context, b *bundle.Bundle) {
 		// Clusters (same as terraform)
 		// https://github.com/databricks/terraform-provider-databricks/blob/v1.75.0/clusters/resource_cluster.go#L315
 		{"resources.clusters.*.autotermination_minutes", 60},
+		{"resources.clusters.*.workload_type.clients.notebooks", true},
+		{"resources.clusters.*.workload_type.clients.jobs", true},
 	}
 
 	for _, defaultDef := range defaults {
diff --git a/bundle/direct/dresources/cluster.go b/bundle/direct/dresources/cluster.go
@@ -134,7 +134,7 @@ func (r *ResourceCluster) ClassifyChange(change structdiff.Change, remoteState *
 }
 
 func makeCreateCluster(config *compute.ClusterSpec) compute.CreateCluster {
-	return compute.CreateCluster{
+	create := compute.CreateCluster{
 		ApplyPolicyDefaultValues:   config.ApplyPolicyDefaultValues,
 		Autoscale:                  config.Autoscale,
 		AutoterminationMinutes:     config.AutoterminationMinutes,
@@ -170,10 +170,18 @@ func makeCreateCluster(config *compute.ClusterSpec) compute.CreateCluster {
 		WorkloadType:               config.WorkloadType,
 		ForceSendFields:            filterFields[compute.CreateCluster](config.ForceSendFields),
 	}
+
+	// If autoscale is not set, we need to send NumWorkers because one of them is required.
+	// If NumWorkers is not nil, we don't need to set it to ForceSendFields as it will be sent anyway.
+	if config.Autoscale == nil && config.NumWorkers == 0 {
+		create.ForceSendFields = append(create.ForceSendFields, "NumWorkers")
+	}
+
+	return create
 }
 
 func makeEditCluster(id string, config *compute.ClusterSpec) compute.EditCluster {
-	return compute.EditCluster{
+	edit := compute.EditCluster{
 		ClusterId:                  id,
 		ApplyPolicyDefaultValues:   config.ApplyPolicyDefaultValues,
 		Autoscale:                  config.Autoscale,
@@ -209,4 +217,12 @@ func makeEditCluster(id string, config *compute.ClusterSpec) compute.EditCluster
 		WorkloadType:               config.WorkloadType,
 		ForceSendFields:            filterFields[compute.EditCluster](config.ForceSendFields),
 	}
+
+	// If autoscale is not set, we need to send NumWorkers because one of them is required.
+	// If NumWorkers is not nil, we don't need to set it to ForceSendFields as it will be sent anyway.
+	if config.Autoscale == nil && config.NumWorkers == 0 {
+		edit.ForceSendFields = append(edit.ForceSendFields, "NumWorkers")
+	}
+
+	return edit
 }

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+trace errcode $CLI bundle deploy`
	`2`	`+trace print_requests.py //clusters/create \| jq -s 'sort_by(.body.cluster_name)[]'`