From bf68318a0cc1932d37ed451919adbfeed9c93622 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Mon, 20 Jan 2025 19:05:16 +0100
Subject: [PATCH] Set bundle auth configuration in command context

---
 .../mutator/initialize_workspace_client.go    | 26 -------------------
 cmd/root/bundle.go                            | 13 ++++++++++
 2 files changed, 13 insertions(+), 26 deletions(-)
 delete mode 100644 bundle/config/mutator/initialize_workspace_client.go

diff --git a/bundle/config/mutator/initialize_workspace_client.go b/bundle/config/mutator/initialize_workspace_client.go
deleted file mode 100644
index 5c905f40c..000000000
--- a/bundle/config/mutator/initialize_workspace_client.go
+++ /dev/null
@@ -1,26 +0,0 @@
-package mutator
-
-import (
-	"context"
-
-	"github.com/databricks/cli/bundle"
-	"github.com/databricks/cli/libs/diag"
-)
-
-type initializeWorkspaceClient struct{}
-
-func InitializeWorkspaceClient() bundle.Mutator {
-	return &initializeWorkspaceClient{}
-}
-
-func (m *initializeWorkspaceClient) Name() string {
-	return "InitializeWorkspaceClient"
-}
-
-// 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 {
-	_, err := b.InitializeWorkspaceClient()
-	return diag.FromErr(err)
-}
diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go
index 8b98f2cf2..36af1c495 100644
--- a/cmd/root/bundle.go
+++ b/cmd/root/bundle.go
@@ -81,6 +81,19 @@ func configureBundle(cmd *cobra.Command, b *bundle.Bundle) (*bundle.Bundle, diag
 
 	// Configure the workspace profile if the flag has been set.
 	diags = diags.Extend(configureProfile(cmd, b))
+	if diags.HasError() {
+		return b, diags
+	}
+
+	// Set the auth configuration in the command context. This can be used
+	// downstream to initialize a API client.
+	client, err := b.InitializeWorkspaceClient()
+	if err != nil {
+		return b, diags.Extend(diag.FromErr(err))
+	}
+	ctx = context.WithValue(ctx, &configUsed, client.Config)
+	cmd.SetContext(ctx)
+
 	return b, diags
 }