Set bundle auth configuration in command context (#2195)

## Changes

This change is required to enable tracking execution time telemetry for
bundle commands. In order to track execution time for the command
generally, we need to have the databricks auth configuration available
at this section of the code:


41bbd89257/cmd/root/root.go (L99)

In order to do this we can rely on the `configUsed` context key.   

Most commands rely on the `root.MustWorkspaceClient` function which
automatically sets the client config in the `configUsed` context key.
Bundle commands, however, do not do so. They instead store their
workspace clients in the `&bundle.Bundle{}` object.

With this PR, the `configUsed` context key will be set for all `bundle`
commands. Functionally nothing changes.

## Tests
Existing tests. Also manually verified that either
`root.MustConfigureBundle` or `utils.ConfigureBundleWithVariables` is
called for all bundle commands (except `bundle init`) thus ensuring this
context key would be set for all bundle commands.

refs for the functions:
1. `root.MustConfigureBundle`:
41bbd89257/cmd/root/bundle.go (L88)
2. `utils.ConfigureBundleWithVariables`:
41bbd89257/cmd/bundle/utils/utils.go (L19)

---------

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>

bundle/bundle.go

@@ -72,6 +72,7 @@ type Bundle struct {
 	// It can be initialized on demand after loading the configuration.
 	clientOnce sync.Once
 	client     *databricks.WorkspaceClient
+	clientErr  error
 
 	// Files that are synced to the workspace.file_path
 	Files []fileset.File
@@ -134,23 +135,25 @@ func TryLoad(ctx context.Context) (*Bundle, error) {
 	return Load(ctx, root)
 }
 
-func (b *Bundle) InitializeWorkspaceClient() (*databricks.WorkspaceClient, error) {
+func (b *Bundle) WorkspaceClientE() (*databricks.WorkspaceClient, error) {
-	client, err := b.Config.Workspace.Client()
+	b.clientOnce.Do(func() {
-	if err != nil {
+		var err error
-		return nil, fmt.Errorf("cannot resolve bundle auth configuration: %w", err)
+		b.client, err = b.Config.Workspace.Client()
-	}
+		if err != nil {
-	return client, nil
+			b.clientErr = fmt.Errorf("cannot resolve bundle auth configuration: %w", err)
+		}
+	})
+
+	return b.client, b.clientErr
 }
 
 func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
-	b.clientOnce.Do(func() {
+	client, err := b.WorkspaceClientE()
-		var err error
+	if err != nil {
-		b.client, err = b.InitializeWorkspaceClient()
+		panic(err)
-		if err != nil {
+	}
-			panic(err)
+
-		}
+	return client
-	})
-	return b.client
 }
 
 // SetWorkpaceClient sets the workspace client for this bundle.

bundle/config/mutator/initialize_workspace_client.go

@@ -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)
-}

bundle/phases/initialize.go

@@ -34,7 +34,6 @@ func Initialize() bundle.Mutator {
 			// If it is an ancestor, this updates all paths to be relative to the sync root path.
 			mutator.SyncInferRoot(),
 
-			mutator.InitializeWorkspaceClient(),
 			mutator.PopulateCurrentUser(),
 			mutator.LoadGitDetails(),
 

cmd/root/auth.go

@@ -209,7 +209,7 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 		if b != nil {
 			ctx = context.WithValue(ctx, &configUsed, b.Config.Workspace.Config())
 			cmd.SetContext(ctx)
-			client, err := b.InitializeWorkspaceClient()
+			client, err := b.WorkspaceClientE()
 			if err != nil {
 				return err
 			}

cmd/root/bundle.go

@@ -81,6 +81,22 @@ 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.
+	//
+	// Note that just initializing a workspace client and loading auth configuration
+	// is a fast operation. It does not perform network I/O or invoke processes (for example the Azure CLI).
+	client, err := b.WorkspaceClientE()
+	if err != nil {
+		return b, diags.Extend(diag.FromErr(err))
+	}
+	ctx = context.WithValue(ctx, &configUsed, client.Config)
+	cmd.SetContext(ctx)
+
 	return b, diags
 }
 

cmd/root/bundle_test.go

@@ -8,7 +8,6 @@ import (
 	"runtime"
 	"testing"
 
-	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
@@ -38,7 +37,7 @@ func emptyCommand(t *testing.T) *cobra.Command {
 	return cmd
 }
 
-func setupWithHost(t *testing.T, cmd *cobra.Command, host string) *bundle.Bundle {
+func setupWithHost(t *testing.T, cmd *cobra.Command, host string) error {
 	setupDatabricksCfg(t)
 
 	rootPath := t.TempDir()
@@ -51,12 +50,11 @@ workspace:
 	err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644)
 	require.NoError(t, err)
 
-	b, diags := MustConfigureBundle(cmd)
+	_, diags := MustConfigureBundle(cmd)
-	require.NoError(t, diags.Error())
+	return diags.Error()
-	return b
 }
 
-func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) *bundle.Bundle {
+func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) error {
 	setupDatabricksCfg(t)
 
 	rootPath := t.TempDir()
@@ -69,29 +67,25 @@ workspace:
 	err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644)
 	require.NoError(t, err)
 
-	b, diags := MustConfigureBundle(cmd)
+	_, diags := MustConfigureBundle(cmd)
-	require.NoError(t, diags.Error())
+	return diags.Error()
-	return b
 }
 
 func TestBundleConfigureDefault(t *testing.T) {
 	testutil.CleanupEnvironment(t)
 
 	cmd := emptyCommand(t)
-	b := setupWithHost(t, cmd, "https://x.com")
+	err := setupWithHost(t, cmd, "https://x.com")
-
-	client, err := b.InitializeWorkspaceClient()
 	require.NoError(t, err)
-	assert.Equal(t, "https://x.com", client.Config.Host)
+
+	assert.Equal(t, "https://x.com", ConfigUsed(cmd.Context()).Host)
 }
 
 func TestBundleConfigureWithMultipleMatches(t *testing.T) {
 	testutil.CleanupEnvironment(t)
 
 	cmd := emptyCommand(t)
-	b := setupWithHost(t, cmd, "https://a.com")
+	err := setupWithHost(t, cmd, "https://a.com")
-
-	_, err := b.InitializeWorkspaceClient()
 	assert.ErrorContains(t, err, "multiple profiles matched: PROFILE-1, PROFILE-2")
 }
 
@@ -101,9 +95,8 @@ func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("NOEXIST")
 	require.NoError(t, err)
-	b := setupWithHost(t, cmd, "https://x.com")
 
-	_, err = b.InitializeWorkspaceClient()
+	err = setupWithHost(t, cmd, "https://x.com")
 	assert.ErrorContains(t, err, "has no NOEXIST profile configured")
 }
 
@@ -113,9 +106,8 @@ func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("PROFILE-1")
 	require.NoError(t, err)
-	b := setupWithHost(t, cmd, "https://x.com")
 
-	_, err = b.InitializeWorkspaceClient()
+	err = setupWithHost(t, cmd, "https://x.com")
 	assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
 }
 
@@ -125,12 +117,11 @@ func TestBundleConfigureWithCorrectProfile(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("PROFILE-1")
 	require.NoError(t, err)
-	b := setupWithHost(t, cmd, "https://a.com")
+	err = setupWithHost(t, cmd, "https://a.com")
 
-	client, err := b.InitializeWorkspaceClient()
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "PROFILE-1", client.Config.Profile)
+	assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
@@ -138,9 +129,8 @@ func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
 
 	t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
 	cmd := emptyCommand(t)
-	b := setupWithHost(t, cmd, "https://x.com")
 
-	_, err := b.InitializeWorkspaceClient()
+	err := setupWithHost(t, cmd, "https://x.com")
 	assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
 }
 
@@ -151,12 +141,11 @@ func TestBundleConfigureWithProfileFlagAndEnvVariable(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("PROFILE-1")
 	require.NoError(t, err)
-	b := setupWithHost(t, cmd, "https://a.com")
 
-	client, err := b.InitializeWorkspaceClient()
+	err = setupWithHost(t, cmd, "https://a.com")
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "PROFILE-1", client.Config.Profile)
+	assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestBundleConfigureProfileDefault(t *testing.T) {
@@ -164,13 +153,12 @@ func TestBundleConfigureProfileDefault(t *testing.T) {
 
 	// The profile in the databricks.yml file is used
 	cmd := emptyCommand(t)
-	b := setupWithProfile(t, cmd, "PROFILE-1")
 
-	client, err := b.InitializeWorkspaceClient()
+	err := setupWithProfile(t, cmd, "PROFILE-1")
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "a", client.Config.Token)
+	assert.Equal(t, "a", ConfigUsed(cmd.Context()).Token)
-	assert.Equal(t, "PROFILE-1", client.Config.Profile)
+	assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestBundleConfigureProfileFlag(t *testing.T) {
@@ -180,13 +168,12 @@ func TestBundleConfigureProfileFlag(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("PROFILE-2")
 	require.NoError(t, err)
-	b := setupWithProfile(t, cmd, "PROFILE-1")
 
-	client, err := b.InitializeWorkspaceClient()
+	err = setupWithProfile(t, cmd, "PROFILE-1")
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
-	assert.Equal(t, "PROFILE-2", client.Config.Profile)
+	assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestBundleConfigureProfileEnvVariable(t *testing.T) {
@@ -195,13 +182,12 @@ func TestBundleConfigureProfileEnvVariable(t *testing.T) {
 	// The DATABRICKS_CONFIG_PROFILE environment variable takes precedence over the profile in the databricks.yml file
 	t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-2")
 	cmd := emptyCommand(t)
-	b := setupWithProfile(t, cmd, "PROFILE-1")
 
-	client, err := b.InitializeWorkspaceClient()
+	err := setupWithProfile(t, cmd, "PROFILE-1")
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
-	assert.Equal(t, "PROFILE-2", client.Config.Profile)
+	assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
@@ -212,13 +198,12 @@ func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
 	cmd := emptyCommand(t)
 	err := cmd.Flag("profile").Value.Set("PROFILE-2")
 	require.NoError(t, err)
-	b := setupWithProfile(t, cmd, "PROFILE-1")
 
-	client, err := b.InitializeWorkspaceClient()
+	err = setupWithProfile(t, cmd, "PROFILE-1")
 	require.NoError(t, err)
-	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
-	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
-	assert.Equal(t, "PROFILE-2", client.Config.Profile)
+	assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
 }
 
 func TestTargetFlagFull(t *testing.T) {