From c07f1370d244bd396cb20905ba269ade4d1e8c38 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 21 Jan 2025 16:43:20 +0100
Subject: [PATCH] Revert "set client during initialization"

This reverts commit 8cc50c2a19ccfc878ef8da96f668240cb537868a.
---
 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, 12 insertions(+), 16 deletions(-)

diff --git a/bundle/bundle.go b/bundle/bundle.go
index 5d19b52e4..3bf4ffb62 100644
--- a/bundle/bundle.go
+++ b/bundle/bundle.go
@@ -143,14 +143,20 @@ func (b *Bundle) InitializeWorkspaceClient() (*databricks.WorkspaceClient, error
 }
 
 func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
-	if b.client == nil {
-		panic("workspace client not initialized yet. This is a bug in the Databricks CLI.")
-	}
+	b.clientOnce.Do(func() {
+		var err error
+		b.client, err = b.InitializeWorkspaceClient()
+		if err != nil {
+			panic(err)
+		}
+	})
 	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 ac48fd484..c7a4ffe4a 100644
--- a/bundle/deploy/terraform/init_test.go
+++ b/bundle/deploy/terraform/init_test.go
@@ -48,10 +48,7 @@ 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")
-
-	client, err := b.InitializeWorkspaceClient()
-	require.NoError(t, err)
-	b.SetWorkpaceClient(client)
+	b.WorkspaceClient()
 
 	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 e353ad48e..b7243ca19 100644
--- a/bundle/deploy/terraform/load_test.go
+++ b/bundle/deploy/terraform/load_test.go
@@ -30,10 +30,7 @@ func TestLoadWithNoState(t *testing.T) {
 
 	t.Setenv("DATABRICKS_HOST", "https://x")
 	t.Setenv("DATABRICKS_TOKEN", "foobar")
-
-	client, err := b.InitializeWorkspaceClient()
-	require.NoError(t, err)
-	b.SetWorkpaceClient(client)
+	b.WorkspaceClient()
 
 	diags := bundle.Apply(context.Background(), b, bundle.Seq(
 		Initialize(),
diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go
index bf20a421f..36af1c495 100644
--- a/cmd/root/bundle.go
+++ b/cmd/root/bundle.go
@@ -93,7 +93,6 @@ 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 d5fe9fffe..b2ec388bd 100644
--- a/libs/template/renderer_test.go
+++ b/libs/template/renderer_test.go
@@ -91,10 +91,7 @@ func assertBuiltinTemplateValid(t *testing.T, template string, settings map[stri
 	})
 
 	b.Tagging = tags.ForCloud(w.Config)
-
-	client, err := b.InitializeWorkspaceClient()
-	require.NoError(t, err)
-	b.SetWorkpaceClient(client)
+	b.WorkspaceClient()
 
 	diags = bundle.Apply(ctx, b, bundle.Seq(
 		phases.Initialize(),