From fc4070994a5b22f9a22557ee4cf80996c8009d86 Mon Sep 17 00:00:00 2001
From: Lennart Kats <lennart.kats@databricks.com>
Date: Fri, 25 Aug 2023 11:01:03 +0200
Subject: [PATCH 1/2] Support cluster overrides with cluster_key and
 compute_key

---
 bundle/config/mutator/override_compute.go      | 10 ++++------
 bundle/config/mutator/override_compute_test.go | 12 ++++++++++++
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/bundle/config/mutator/override_compute.go b/bundle/config/mutator/override_compute.go
index 124392491d..9e826a0f3f 100644
--- a/bundle/config/mutator/override_compute.go
+++ b/bundle/config/mutator/override_compute.go
@@ -23,12 +23,10 @@ func (m *overrideCompute) Name() string {
 func overrideJobCompute(j *resources.Job, compute string) {
 	for i := range j.Tasks {
 		task := &j.Tasks[i]
-		if task.NewCluster != nil {
-			task.NewCluster = nil
-			task.ExistingClusterId = compute
-		} else if task.ExistingClusterId != "" {
-			task.ExistingClusterId = compute
-		}
+		task.NewCluster = nil
+		task.JobClusterKey = ""
+		task.ComputeKey = ""
+		task.ExistingClusterId = compute
 	}
 }
 
diff --git a/bundle/config/mutator/override_compute_test.go b/bundle/config/mutator/override_compute_test.go
index 9eb99edb9b..994c95da73 100644
--- a/bundle/config/mutator/override_compute_test.go
+++ b/bundle/config/mutator/override_compute_test.go
@@ -34,6 +34,12 @@ func TestOverrideDevelopment(t *testing.T) {
 							{
 								ExistingClusterId: "cluster2",
 							},
+							{
+								ComputeKey: "compute_key",
+							},
+							{
+								JobClusterKey: "cluster_key",
+							},
 						},
 					}},
 				},
@@ -47,6 +53,12 @@ func TestOverrideDevelopment(t *testing.T) {
 	assert.Nil(t, bundle.Config.Resources.Jobs["job1"].Tasks[0].NewCluster)
 	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[0].ExistingClusterId)
 	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[1].ExistingClusterId)
+	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[2].ExistingClusterId)
+	assert.Equal(t, "newClusterID", bundle.Config.Resources.Jobs["job1"].Tasks[3].ExistingClusterId)
+
+	assert.Nil(t, bundle.Config.Resources.Jobs["job1"].Tasks[0].NewCluster)
+	assert.Empty(t, bundle.Config.Resources.Jobs["job1"].Tasks[2].ComputeKey)
+	assert.Empty(t, bundle.Config.Resources.Jobs["job1"].Tasks[3].JobClusterKey)
 }
 
 func TestOverrideDevelopmentEnv(t *testing.T) {

From 3b740a210fe80bf5ce3861def99b3f5d1abf5b06 Mon Sep 17 00:00:00 2001
From: Lennart Kats <lennart.kats@databricks.com>
Date: Fri, 25 Aug 2023 13:41:05 +0200
Subject: [PATCH 2/2] Fix overriding compute for pipeline tasks

---
 bundle/config/mutator/override_compute.go     | 10 +++++---
 .../config/mutator/override_compute_test.go   | 25 +++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/bundle/config/mutator/override_compute.go b/bundle/config/mutator/override_compute.go
index 9e826a0f3f..ee2e2a8254 100644
--- a/bundle/config/mutator/override_compute.go
+++ b/bundle/config/mutator/override_compute.go
@@ -23,10 +23,12 @@ func (m *overrideCompute) Name() string {
 func overrideJobCompute(j *resources.Job, compute string) {
 	for i := range j.Tasks {
 		task := &j.Tasks[i]
-		task.NewCluster = nil
-		task.JobClusterKey = ""
-		task.ComputeKey = ""
-		task.ExistingClusterId = compute
+		if task.NewCluster != nil || task.ExistingClusterId != "" || task.ComputeKey != "" || task.JobClusterKey != "" {
+			task.NewCluster = nil
+			task.JobClusterKey = ""
+			task.ComputeKey = ""
+			task.ExistingClusterId = compute
+		}
 	}
 }
 
diff --git a/bundle/config/mutator/override_compute_test.go b/bundle/config/mutator/override_compute_test.go
index 994c95da73..f04c91c460 100644
--- a/bundle/config/mutator/override_compute_test.go
+++ b/bundle/config/mutator/override_compute_test.go
@@ -89,6 +89,31 @@ func TestOverrideDevelopmentEnv(t *testing.T) {
 	assert.Equal(t, "cluster2", bundle.Config.Resources.Jobs["job1"].Tasks[1].ExistingClusterId)
 }
 
+func TestOverridePipelineTask(t *testing.T) {
+	os.Setenv("DATABRICKS_CLUSTER_ID", "newClusterId")
+	bundle := &bundle.Bundle{
+		Config: config.Root{
+			Resources: config.Resources{
+				Jobs: map[string]*resources.Job{
+					"job1": {JobSettings: &jobs.JobSettings{
+						Name: "job1",
+						Tasks: []jobs.Task{
+							{
+								PipelineTask: &jobs.PipelineTask{},
+							},
+						},
+					}},
+				},
+			},
+		},
+	}
+
+	m := mutator.OverrideCompute()
+	err := m.Apply(context.Background(), bundle)
+	require.NoError(t, err)
+	assert.Empty(t, bundle.Config.Resources.Jobs["job1"].Tasks[0].ExistingClusterId)
+}
+
 func TestOverrideProduction(t *testing.T) {
 	bundle := &bundle.Bundle{
 		Config: config.Root{