From d221d4d34811dee157a7da82a8fcd3e16da9cd54 Mon Sep 17 00:00:00 2001
From: Lennart Kats <lennart.kats@databricks.com>
Date: Fri, 7 Jul 2023 09:58:18 +0200
Subject: [PATCH] Make it possible to override the compute id in an environment

---
 bundle/config/bundle.go                        | 4 ++--
 bundle/config/environment.go                   | 6 ++++++
 bundle/config/mutator/override_compute.go      | 9 +++++----
 bundle/config/mutator/override_compute_test.go | 6 +++---
 bundle/config/root.go                          | 4 ++++
 cmd/bundle/deploy.go                           | 4 ++--
 6 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/bundle/config/bundle.go b/bundle/config/bundle.go
index 49c8626da..382e89be3 100644
--- a/bundle/config/bundle.go
+++ b/bundle/config/bundle.go
@@ -32,8 +32,8 @@ type Bundle struct {
 	// Determines the mode of the environment.
 	// For example, 'mode: development' can be used for deployments for
 	// development purposes.
-	Mode Mode `json:"mode,omitempty"`
+	Mode Mode `json:"mode,omitempty" bundle:"readonly"`
 
 	// Overrides the compute used for jobs and other supported assets.
-	Compute string `json:"override_compute,omitempty" bundle:"readonly"`
+	ComputeID string `json:"compute_id,omitempty"`
 }
diff --git a/bundle/config/environment.go b/bundle/config/environment.go
index 5e58b9a78..06a8d8909 100644
--- a/bundle/config/environment.go
+++ b/bundle/config/environment.go
@@ -9,8 +9,14 @@ type Environment struct {
 	// by the user (through environment variable or command line argument).
 	Default bool `json:"default,omitempty"`
 
+	// Determines the mode of the environment.
+	// For example, 'mode: development' can be used for deployments for
+	// development purposes.
 	Mode Mode `json:"mode,omitempty"`
 
+	// Overrides the compute used for jobs and other supported assets.
+	ComputeID string `json:"compute_id,omitempty"`
+
 	Bundle *Bundle `json:"bundle,omitempty"`
 
 	Workspace *Workspace `json:"workspace,omitempty"`
diff --git a/bundle/config/mutator/override_compute.go b/bundle/config/mutator/override_compute.go
index 8ce28cbcc..ba3fd9940 100644
--- a/bundle/config/mutator/override_compute.go
+++ b/bundle/config/mutator/override_compute.go
@@ -34,21 +34,22 @@ func overrideJobCompute(j *resources.Job, compute string) {
 
 func (m *overrideCompute) Apply(ctx context.Context, b *bundle.Bundle) error {
 	if b.Config.Bundle.Mode != config.Development {
-		if b.Config.Bundle.Compute != "" {
+		if b.Config.Bundle.ComputeID != "" {
 			return fmt.Errorf("cannot override compute for an environment that does not use 'mode: development'")
 		}
 		return nil
 	}
 	if os.Getenv("DATABRICKS_CLUSTER_ID") != "" {
-		b.Config.Bundle.Compute = os.Getenv("DATABRICKS_CLUSTER_ID")
+		b.Config.Bundle.ComputeID = os.Getenv("DATABRICKS_CLUSTER_ID")
 	}
-	if b.Config.Bundle.Compute == "" {
+
+	if b.Config.Bundle.ComputeID == "" {
 		return nil
 	}
 
 	r := b.Config.Resources
 	for i := range r.Jobs {
-		overrideJobCompute(r.Jobs[i], b.Config.Bundle.Compute)
+		overrideJobCompute(r.Jobs[i], b.Config.Bundle.ComputeID)
 	}
 
 	return nil
diff --git a/bundle/config/mutator/override_compute_test.go b/bundle/config/mutator/override_compute_test.go
index a9da6856f..9eb99edb9 100644
--- a/bundle/config/mutator/override_compute_test.go
+++ b/bundle/config/mutator/override_compute_test.go
@@ -20,8 +20,8 @@ func TestOverrideDevelopment(t *testing.T) {
 	bundle := &bundle.Bundle{
 		Config: config.Root{
 			Bundle: config.Bundle{
-				Mode:    config.Development,
-				Compute: "newClusterID",
+				Mode:      config.Development,
+				ComputeID: "newClusterID",
 			},
 			Resources: config.Resources{
 				Jobs: map[string]*resources.Job{
@@ -81,7 +81,7 @@ func TestOverrideProduction(t *testing.T) {
 	bundle := &bundle.Bundle{
 		Config: config.Root{
 			Bundle: config.Bundle{
-				Compute: "newClusterID",
+				ComputeID: "newClusterID",
 			},
 			Resources: config.Resources{
 				Jobs: map[string]*resources.Job{
diff --git a/bundle/config/root.go b/bundle/config/root.go
index ac3aa20da..f1ae9c0a5 100644
--- a/bundle/config/root.go
+++ b/bundle/config/root.go
@@ -194,5 +194,9 @@ func (r *Root) MergeEnvironment(env *Environment) error {
 		r.Bundle.Mode = env.Mode
 	}
 
+	if env.ComputeID != "" {
+		r.Bundle.ComputeID = env.ComputeID
+	}
+
 	return nil
 }
diff --git a/cmd/bundle/deploy.go b/cmd/bundle/deploy.go
index 05df84c59..e87570199 100644
--- a/cmd/bundle/deploy.go
+++ b/cmd/bundle/deploy.go
@@ -21,7 +21,7 @@ var deployCmd = &cobra.Command{
 func deploy(cmd *cobra.Command, b *bundle.Bundle) error {
 	// If `--force` is specified, force acquisition of the deployment lock.
 	b.Config.Bundle.Lock.Force = forceDeploy
-	b.Config.Bundle.Compute = computeID
+	b.Config.Bundle.ComputeID = computeID
 
 	return bundle.Apply(cmd.Context(), b, bundle.Seq(
 		phases.Initialize(),
@@ -36,5 +36,5 @@ var computeID string
 func init() {
 	AddCommand(deployCmd)
 	deployCmd.Flags().BoolVar(&forceDeploy, "force", false, "Force acquisition of deployment lock.")
-	deployCmd.Flags().StringVar(&computeID, "compute", "", "Override compute in the deployment with the given compute ID.")
+	deployCmd.Flags().StringVarP(&computeID, "compute-id", "c", "", "Override compute in the deployment with the given compute ID.")
 }