From 861f33d37696413f07fea6c6bb72fa1f8486fda6 Mon Sep 17 00:00:00 2001
From: "Lennart Kats (databricks)" <lennart.kats@databricks.com>
Date: Mon, 28 Aug 2023 00:51:35 -0700
Subject: [PATCH] Support cluster overrides with cluster_key and compute_key
 (#696)

## Changes

Support `databricks bundle deploy --compute-id my_all_purpose_cluster`
in two missing cases.
---
 bundle/config/mutator/override_compute.go     |  6 +--
 .../config/mutator/override_compute_test.go   | 37 +++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/bundle/config/mutator/override_compute.go b/bundle/config/mutator/override_compute.go
index 124392491..ee2e2a825 100644
--- a/bundle/config/mutator/override_compute.go
+++ b/bundle/config/mutator/override_compute.go
@@ -23,10 +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 {
+		if task.NewCluster != nil || task.ExistingClusterId != "" || task.ComputeKey != "" || task.JobClusterKey != "" {
 			task.NewCluster = nil
-			task.ExistingClusterId = compute
-		} else if task.ExistingClusterId != "" {
+			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 9eb99edb9..f04c91c46 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) {
@@ -77,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{