From e403630e4dfecebabf8124c552e18e1b6bc88acb Mon Sep 17 00:00:00 2001
From: Anders Rex <anders.rex@databricks.com>
Date: Thu, 18 Jul 2024 13:39:19 +0300
Subject: [PATCH] Check for host mismatch when initializing workspace client

---
 .../mutator/initialize_workspace_client.go    | 18 ++++++++++-
 ...go => initialize_workspace_client_test.go} |  8 ++---
 .../config/mutator/validate_workspace_host.go | 30 -------------------
 bundle/phases/deploy.go                       |  1 -
 4 files changed, 21 insertions(+), 36 deletions(-)
 rename bundle/config/mutator/{validate_workspace_host_test.go => initialize_workspace_client_test.go} (81%)
 delete mode 100644 bundle/config/mutator/validate_workspace_host.go

diff --git a/bundle/config/mutator/initialize_workspace_client.go b/bundle/config/mutator/initialize_workspace_client.go
index 5c905f40..ea2f1d06 100644
--- a/bundle/config/mutator/initialize_workspace_client.go
+++ b/bundle/config/mutator/initialize_workspace_client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/env"
 	"github.com/databricks/cli/libs/diag"
 )
 
@@ -20,7 +21,22 @@ func (m *initializeWorkspaceClient) Name() string {
 // Apply initializes the workspace client for the bundle. We do this here so
 // downstream calls to b.WorkspaceClient() do not panic if there's an error in the
 // auth configuration.
-func (m *initializeWorkspaceClient) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
+func (m *initializeWorkspaceClient) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	if err := validateWorkspaceHost(ctx, b); err != nil {
+		return err
+	}
+
 	_, err := b.InitializeWorkspaceClient()
 	return diag.FromErr(err)
 }
+
+func validateWorkspaceHost(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	env_host, _ := env.Host(ctx)
+	target_host := b.Config.Workspace.Host
+
+	if env_host != "" && target_host != "" && env_host != target_host {
+		return diag.Errorf("target host and DATABRICKS_HOST environment variable mismatch")
+	}
+
+	return nil
+}
diff --git a/bundle/config/mutator/validate_workspace_host_test.go b/bundle/config/mutator/initialize_workspace_client_test.go
similarity index 81%
rename from bundle/config/mutator/validate_workspace_host_test.go
rename to bundle/config/mutator/initialize_workspace_client_test.go
index 718d2dac..1729f030 100644
--- a/bundle/config/mutator/validate_workspace_host_test.go
+++ b/bundle/config/mutator/initialize_workspace_client_test.go
@@ -10,7 +10,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestValidateValidateWorkspaceHost(t *testing.T) {
+func TestInitializeWorkspaceClient(t *testing.T) {
 	host := "https://host.databricks.com"
 	b := &bundle.Bundle{
 		Config: config.Root{
@@ -21,12 +21,12 @@ func TestValidateValidateWorkspaceHost(t *testing.T) {
 	}
 	t.Setenv(env.HostVariable, host)
 
-	m := ValidateWorkspaceHost()
+	m := InitializeWorkspaceClient()
 	diags := bundle.Apply(context.Background(), b, m)
 	assert.NoError(t, diags.Error())
 }
 
-func TestValidateValidateWorkspaceHostMismatch(t *testing.T) {
+func TestInitializeWorkspaceClientWorkspaceHostMismatch(t *testing.T) {
 	b := &bundle.Bundle{
 		Config: config.Root{
 			Workspace: config.Workspace{
@@ -36,7 +36,7 @@ func TestValidateValidateWorkspaceHostMismatch(t *testing.T) {
 	}
 	t.Setenv(env.HostVariable, "https://env-host.databricks.com")
 
-	m := ValidateWorkspaceHost()
+	m := InitializeWorkspaceClient()
 	diags := bundle.Apply(context.Background(), b, m)
 	expectedError := "target host and DATABRICKS_HOST environment variable mismatch"
 	assert.EqualError(t, diags.Error(), expectedError)
diff --git a/bundle/config/mutator/validate_workspace_host.go b/bundle/config/mutator/validate_workspace_host.go
deleted file mode 100644
index 5450ab96..00000000
--- a/bundle/config/mutator/validate_workspace_host.go
+++ /dev/null
@@ -1,30 +0,0 @@
-package mutator
-
-import (
-	"context"
-
-	"github.com/databricks/cli/bundle"
-	"github.com/databricks/cli/bundle/env"
-	"github.com/databricks/cli/libs/diag"
-)
-
-type validateWorkspaceHost struct{}
-
-func ValidateWorkspaceHost() *validateWorkspaceHost {
-	return &validateWorkspaceHost{}
-}
-
-func (m *validateWorkspaceHost) Name() string {
-	return "ValidateWorkspaceHost"
-}
-
-func (m *validateWorkspaceHost) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-	env_host, _ := env.Host(ctx)
-	target_host := b.Config.Workspace.Host
-
-	if env_host != "" && target_host != "" && env_host != target_host {
-		return diag.Errorf("target host and DATABRICKS_HOST environment variable mismatch")
-	}
-
-	return nil
-}
diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go
index e30a69c9..46c38918 100644
--- a/bundle/phases/deploy.go
+++ b/bundle/phases/deploy.go
@@ -25,7 +25,6 @@ func Deploy() bundle.Mutator {
 			bundle.Seq(
 				terraform.StatePull(),
 				deploy.StatePull(),
-				mutator.ValidateWorkspaceHost(),
 				mutator.ValidateGitDetails(),
 				libraries.ValidateLocalLibrariesExist(),
 				artifacts.CleanUp(),