From 8cc50c2a19ccfc878ef8da96f668240cb537868a Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 21 Jan 2025 16:28:25 +0100
Subject: [PATCH] set client during initialization

---
 bundle/bundle.go                     | 12 +++---------
 bundle/deploy/terraform/init_test.go |  5 ++++-
 bundle/deploy/terraform/load_test.go |  5 ++++-
 cmd/root/bundle.go                   |  1 +
 libs/template/renderer_test.go       |  5 ++++-
 5 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/bundle/bundle.go b/bundle/bundle.go
index 3bf4ffb62..5d19b52e4 100644
--- a/bundle/bundle.go
+++ b/bundle/bundle.go
@@ -143,20 +143,14 @@ func (b *Bundle) InitializeWorkspaceClient() (*databricks.WorkspaceClient, error
 }
 
 func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
-	b.clientOnce.Do(func() {
-		var err error
-		b.client, err = b.InitializeWorkspaceClient()
-		if err != nil {
-			panic(err)
-		}
-	})
+	if b.client == nil {
+		panic("workspace client not initialized yet. This is a bug in the Databricks CLI.")
+	}
 	return b.client
 }
 
 // SetWorkpaceClient sets the workspace client for this bundle.
-// This is used to inject a mock client for testing.
 func (b *Bundle) SetWorkpaceClient(w *databricks.WorkspaceClient) {
-	b.clientOnce.Do(func() {})
 	b.client = w
 }
 
diff --git a/bundle/deploy/terraform/init_test.go b/bundle/deploy/terraform/init_test.go
index c7a4ffe4a..ac48fd484 100644
--- a/bundle/deploy/terraform/init_test.go
+++ b/bundle/deploy/terraform/init_test.go
@@ -48,7 +48,10 @@ func TestInitEnvironmentVariables(t *testing.T) {
 	// TODO(pietern): create test fixture that initializes a mocked client.
 	t.Setenv("DATABRICKS_HOST", "https://x")
 	t.Setenv("DATABRICKS_TOKEN", "foobar")
-	b.WorkspaceClient()
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	b.SetWorkpaceClient(client)
 
 	diags := bundle.Apply(context.Background(), b, Initialize())
 	require.NoError(t, diags.Error())
diff --git a/bundle/deploy/terraform/load_test.go b/bundle/deploy/terraform/load_test.go
index b7243ca19..e353ad48e 100644
--- a/bundle/deploy/terraform/load_test.go
+++ b/bundle/deploy/terraform/load_test.go
@@ -30,7 +30,10 @@ func TestLoadWithNoState(t *testing.T) {
 
 	t.Setenv("DATABRICKS_HOST", "https://x")
 	t.Setenv("DATABRICKS_TOKEN", "foobar")
-	b.WorkspaceClient()
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	b.SetWorkpaceClient(client)
 
 	diags := bundle.Apply(context.Background(), b, bundle.Seq(
 		Initialize(),
diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go
index 36af1c495..bf20a421f 100644
--- a/cmd/root/bundle.go
+++ b/cmd/root/bundle.go
@@ -93,6 +93,7 @@ func configureBundle(cmd *cobra.Command, b *bundle.Bundle) (*bundle.Bundle, diag
 	}
 	ctx = context.WithValue(ctx, &configUsed, client.Config)
 	cmd.SetContext(ctx)
+	b.SetWorkpaceClient(client)
 
 	return b, diags
 }
diff --git a/libs/template/renderer_test.go b/libs/template/renderer_test.go
index b2ec388bd..d5fe9fffe 100644
--- a/libs/template/renderer_test.go
+++ b/libs/template/renderer_test.go
@@ -91,7 +91,10 @@ func assertBuiltinTemplateValid(t *testing.T, template string, settings map[stri
 	})
 
 	b.Tagging = tags.ForCloud(w.Config)
-	b.WorkspaceClient()
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	b.SetWorkpaceClient(client)
 
 	diags = bundle.Apply(ctx, b, bundle.Seq(
 		phases.Initialize(),