Make bundle loaders return diagnostics (#1319)

## Changes

The function signature of Cobra's `PreRunE` function has an `error`
return value. We'd like to start returning `diag.Diagnostics` after
loading a bundle, so this is incompatible. This change modifies all
usage of `PreRunE` to load a bundle to inline function calls in the
command's `RunE` function.

## Tests

* Unit tests pass.
* Integration tests pass.

cmd/bundle/deploy.go

@@ -13,10 +13,9 @@ import (
 
 func newDeployCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "deploy",
-		Short:   "Deploy bundle",
-		Args:    root.NoArgs,
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "deploy",
+		Short: "Deploy bundle",
+		Args:  root.NoArgs,
 	}
 
 	var force bool
@@ -30,7 +29,10 @@ func newDeployCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
 		bundle.ApplyFunc(ctx, b, func(context.Context, *bundle.Bundle) diag.Diagnostics {
 			b.Config.Bundle.Force = force
@@ -46,7 +48,7 @@ func newDeployCommand() *cobra.Command {
 			return nil
 		})
 
-		diags := bundle.Apply(ctx, b, bundle.Seq(
+		diags = bundle.Apply(ctx, b, bundle.Seq(
 			phases.Initialize(),
 			phases.Build(),
 			phases.Deploy(),

cmd/bundle/deployment/bind.go

@@ -16,10 +16,9 @@ import (
 
 func newBindCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "bind KEY RESOURCE_ID",
-		Short:   "Bind bundle-defined resources to existing resources",
-		Args:    root.ExactArgs(2),
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "bind KEY RESOURCE_ID",
+		Short: "Bind bundle-defined resources to existing resources",
+		Args:  root.ExactArgs(2),
 	}
 
 	var autoApprove bool
@@ -29,7 +28,11 @@ func newBindCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
+
 		resource, err := b.Config.Resources.FindResourceByConfigKey(args[0])
 		if err != nil {
 			return err
@@ -50,7 +53,7 @@ func newBindCommand() *cobra.Command {
 			return nil
 		})
 
-		diags := bundle.Apply(ctx, b, bundle.Seq(
+		diags = bundle.Apply(ctx, b, bundle.Seq(
 			phases.Initialize(),
 			phases.Bind(&terraform.BindOptions{
 				AutoApprove:  autoApprove,

cmd/bundle/deployment/unbind.go

@@ -13,10 +13,9 @@ import (
 
 func newUnbindCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "unbind KEY",
-		Short:   "Unbind bundle-defined resources from its managed remote resource",
-		Args:    root.ExactArgs(1),
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "unbind KEY",
+		Short: "Unbind bundle-defined resources from its managed remote resource",
+		Args:  root.ExactArgs(1),
 	}
 
 	var forceLock bool
@@ -24,7 +23,11 @@ func newUnbindCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
+
 		resource, err := b.Config.Resources.FindResourceByConfigKey(args[0])
 		if err != nil {
 			return err
@@ -35,7 +38,7 @@ func newUnbindCommand() *cobra.Command {
 			return nil
 		})
 
-		diags := bundle.Apply(cmd.Context(), b, bundle.Seq(
+		diags = bundle.Apply(cmd.Context(), b, bundle.Seq(
 			phases.Initialize(),
 			phases.Unbind(resource.TerraformResourceName(), args[0]),
 		))

cmd/bundle/destroy.go

@@ -18,10 +18,9 @@ import (
 
 func newDestroyCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "destroy",
-		Short:   "Destroy deployed bundle resources",
-		Args:    root.NoArgs,
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "destroy",
+		Short: "Destroy deployed bundle resources",
+		Args:  root.NoArgs,
 	}
 
 	var autoApprove bool
@@ -31,7 +30,10 @@ func newDestroyCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
 		bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 			// If `--force-lock` is specified, force acquisition of the deployment lock.
@@ -58,7 +60,7 @@ func newDestroyCommand() *cobra.Command {
 			return fmt.Errorf("please specify --auto-approve since selected logging format is json")
 		}
 
-		diags := bundle.Apply(ctx, b, bundle.Seq(
+		diags = bundle.Apply(ctx, b, bundle.Seq(
 			phases.Initialize(),
 			phases.Build(),
 			phases.Destroy(),

cmd/bundle/generate.go

@@ -2,7 +2,6 @@ package bundle
 
 import (
 	"github.com/databricks/cli/cmd/bundle/generate"
-	"github.com/databricks/cli/cmd/bundle/utils"
 	"github.com/spf13/cobra"
 )
 
@@ -10,10 +9,9 @@ func newGenerateCommand() *cobra.Command {
 	var key string
 
 	cmd := &cobra.Command{
-		Use:     "generate",
-		Short:   "Generate bundle configuration",
-		Long:    "Generate bundle configuration",
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "generate",
+		Short: "Generate bundle configuration",
+		Long:  "Generate bundle configuration",
 	}
 
 	cmd.AddCommand(generate.NewGenerateJobCommand())

cmd/bundle/generate/job.go

@@ -5,7 +5,6 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/generate"
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
@@ -24,9 +23,8 @@ func NewGenerateJobCommand() *cobra.Command {
 	var force bool
 
 	cmd := &cobra.Command{
-		Use:     "job",
-		Short:   "Generate bundle configuration for a job",
-		PreRunE: root.MustConfigureBundle,
+		Use:   "job",
+		Short: "Generate bundle configuration for a job",
 	}
 
 	cmd.Flags().Int64Var(&jobId, "existing-job-id", 0, `Job ID of the job to generate config for`)
@@ -43,9 +41,12 @@ func NewGenerateJobCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
-		w := b.WorkspaceClient()
+		b, diags := root.MustConfigureBundle(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
+		w := b.WorkspaceClient()
 		job, err := w.Jobs.Get(ctx, jobs.GetJobRequest{JobId: jobId})
 		if err != nil {
 			return err

cmd/bundle/generate/pipeline.go

@@ -5,7 +5,6 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config/generate"
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
@@ -24,9 +23,8 @@ func NewGeneratePipelineCommand() *cobra.Command {
 	var force bool
 
 	cmd := &cobra.Command{
-		Use:     "pipeline",
-		Short:   "Generate bundle configuration for a pipeline",
-		PreRunE: root.MustConfigureBundle,
+		Use:   "pipeline",
+		Short: "Generate bundle configuration for a pipeline",
 	}
 
 	cmd.Flags().StringVar(&pipelineId, "existing-pipeline-id", "", `ID of the pipeline to generate config for`)
@@ -43,9 +41,12 @@ func NewGeneratePipelineCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
-		w := b.WorkspaceClient()
+		b, diags := root.MustConfigureBundle(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
+		w := b.WorkspaceClient()
 		pipeline, err := w.Pipelines.Get(ctx, pipelines.GetPipelineRequest{PipelineId: pipelineId})
 		if err != nil {
 			return err

cmd/bundle/launch.go

@@ -16,8 +16,6 @@ func newLaunchCommand() *cobra.Command {
 
 		// We're not ready to expose this command until we specify its semantics.
 		Hidden: true,
-
-		PreRunE: root.MustConfigureBundle,
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {

cmd/bundle/run.go

@@ -17,10 +17,9 @@ import (
 
 func newRunCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "run [flags] KEY",
-		Short:   "Run a resource (e.g. a job or a pipeline)",
-		Args:    root.MaximumNArgs(1),
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "run [flags] KEY",
+		Short: "Run a resource (e.g. a job or a pipeline)",
+		Args:  root.MaximumNArgs(1),
 	}
 
 	var runOptions run.Options
@@ -33,9 +32,12 @@ func newRunCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		b := bundle.Get(ctx)
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
-		diags := bundle.Apply(ctx, b, bundle.Seq(
+		diags = bundle.Apply(ctx, b, bundle.Seq(
 			phases.Initialize(),
 			terraform.Interpolate(),
 			terraform.Write(),
@@ -109,15 +111,14 @@ func newRunCommand() *cobra.Command {
 			return nil, cobra.ShellCompDirectiveNoFileComp
 		}
 
-		err := root.MustConfigureBundle(cmd, args)
-		if err != nil {
+		b, diags := root.MustConfigureBundle(cmd)
+		if err := diags.Error(); err != nil {
 			cobra.CompErrorln(err.Error())
 			return nil, cobra.ShellCompDirectiveError
 		}
 
 		// No completion in the context of a bundle.
 		// Source and destination paths are taken from bundle configuration.
-		b := bundle.GetOrNil(cmd.Context())
 		if b == nil {
 			return nil, cobra.ShellCompDirectiveNoFileComp
 		}

cmd/bundle/summary.go

@@ -18,10 +18,9 @@ import (
 
 func newSummaryCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "summary",
-		Short:   "Describe the bundle resources and their deployment states",
-		Args:    root.NoArgs,
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "summary",
+		Short: "Describe the bundle resources and their deployment states",
+		Args:  root.NoArgs,
 
 		// This command is currently intended for the Databricks VSCode extension only
 		Hidden: true,
@@ -31,14 +30,18 @@ func newSummaryCommand() *cobra.Command {
 	cmd.Flags().BoolVar(&forcePull, "force-pull", false, "Skip local cache and load the state from the remote workspace")
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		b := bundle.Get(cmd.Context())
+		ctx := cmd.Context()
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
-		diags := bundle.Apply(cmd.Context(), b, phases.Initialize())
+		diags = bundle.Apply(ctx, b, phases.Initialize())
 		if err := diags.Error(); err != nil {
 			return err
 		}
 
-		cacheDir, err := terraform.Dir(cmd.Context(), b)
+		cacheDir, err := terraform.Dir(ctx, b)
 		if err != nil {
 			return err
 		}
@@ -47,7 +50,7 @@ func newSummaryCommand() *cobra.Command {
 		noCache := errors.Is(stateFileErr, os.ErrNotExist) || errors.Is(configFileErr, os.ErrNotExist)
 
 		if forcePull || noCache {
-			diags = bundle.Apply(cmd.Context(), b, bundle.Seq(
+			diags = bundle.Apply(ctx, b, bundle.Seq(
 				terraform.StatePull(),
 				terraform.Interpolate(),
 				terraform.Write(),
@@ -57,7 +60,7 @@ func newSummaryCommand() *cobra.Command {
 			}
 		}
 
-		diags = bundle.Apply(cmd.Context(), b, terraform.Load())
+		diags = bundle.Apply(ctx, b, terraform.Load())
 		if err := diags.Error(); err != nil {
 			return err
 		}

cmd/bundle/sync.go

@@ -36,8 +36,6 @@ func newSyncCommand() *cobra.Command {
 		Use:   "sync [flags]",
 		Short: "Synchronize bundle tree to the workspace",
 		Args:  root.NoArgs,
-
-		PreRunE: utils.ConfigureBundleWithVariables,
 	}
 
 	var f syncFlags
@@ -46,10 +44,14 @@ func newSyncCommand() *cobra.Command {
 	cmd.Flags().BoolVar(&f.watch, "watch", false, "watch local file system for changes")
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		b := bundle.Get(cmd.Context())
+		ctx := cmd.Context()
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
 		// Run initialize phase to make sure paths are set.
-		diags := bundle.Apply(cmd.Context(), b, phases.Initialize())
+		diags = bundle.Apply(ctx, b, phases.Initialize())
 		if err := diags.Error(); err != nil {
 			return err
 		}
@@ -59,7 +61,6 @@ func newSyncCommand() *cobra.Command {
 			return err
 		}
 
-		ctx := cmd.Context()
 		s, err := sync.New(ctx, *opts)
 		if err != nil {
 			return err

cmd/bundle/test.go

@@ -3,7 +3,6 @@ package bundle
 import (
 	"fmt"
 
-	"github.com/databricks/cli/cmd/root"
 	"github.com/spf13/cobra"
 )
 
@@ -15,8 +14,6 @@ func newTestCommand() *cobra.Command {
 
 		// We're not ready to expose this command until we specify its semantics.
 		Hidden: true,
-
-		PreRunE: root.MustConfigureBundle,
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {

cmd/bundle/utils/utils.go

@@ -9,23 +9,30 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func ConfigureBundleWithVariables(cmd *cobra.Command, args []string) error {
+func configureVariables(cmd *cobra.Command, b *bundle.Bundle, variables []string) diag.Diagnostics {
+	return bundle.ApplyFunc(cmd.Context(), b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		err := b.Config.InitializeVariables(variables)
+		return diag.FromErr(err)
+	})
+}
+
+func ConfigureBundleWithVariables(cmd *cobra.Command) (*bundle.Bundle, diag.Diagnostics) {
 	// Load bundle config and apply target
-	err := root.MustConfigureBundle(cmd, args)
-	if err != nil {
-		return err
+	b, diags := root.MustConfigureBundle(cmd)
+	if diags.HasError() {
+		return nil, diags
 	}
 
 	variables, err := cmd.Flags().GetStringSlice("var")
 	if err != nil {
-		return err
+		return nil, diag.FromErr(err)
 	}
 
 	// Initialize variables by assigning them values passed as command line flags
-	b := bundle.Get(cmd.Context())
-	diags := bundle.ApplyFunc(cmd.Context(), b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-		err := b.Config.InitializeVariables(variables)
-		return diag.FromErr(err)
-	})
-	return diags.Error()
+	diags = diags.Extend(configureVariables(cmd, b, variables))
+	if diags.HasError() {
+		return nil, diags
+	}
+
+	return b, diags
 }

cmd/bundle/validate.go

@@ -13,16 +13,19 @@ import (
 
 func newValidateCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "validate",
-		Short:   "Validate configuration",
-		Args:    root.NoArgs,
-		PreRunE: utils.ConfigureBundleWithVariables,
+		Use:   "validate",
+		Short: "Validate configuration",
+		Args:  root.NoArgs,
 	}
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		b := bundle.Get(cmd.Context())
+		ctx := cmd.Context()
+		b, diags := utils.ConfigureBundleWithVariables(cmd)
+		if err := diags.Error(); err != nil {
+			return diags.Error()
+		}
 
-		diags := bundle.Apply(cmd.Context(), b, phases.Initialize())
+		diags = bundle.Apply(ctx, b, phases.Initialize())
 		if err := diags.Error(); err != nil {
 			return err
 		}

cmd/labs/project/entrypoint.go

@@ -10,7 +10,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/cmdio"
@@ -203,11 +202,11 @@ func (e *Entrypoint) getLoginConfig(cmd *cobra.Command) (*loginConfig, *config.C
 		return lc, cfg, nil
 	}
 	if e.IsBundleAware {
-		err = root.TryConfigureBundle(cmd, []string{})
-		if err != nil {
+		b, diags := root.TryConfigureBundle(cmd)
+		if err := diags.Error(); err != nil {
 			return nil, nil, fmt.Errorf("bundle: %w", err)
 		}
-		if b := bundle.GetOrNil(cmd.Context()); b != nil {
+		if b != nil {
 			log.Infof(ctx, "Using login configuration from Databricks Asset Bundle")
 			return &loginConfig{}, b.WorkspaceClient().Config, nil
 		}

cmd/root/auth.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"net/http"
 
-	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/databricks-sdk-go"
@@ -149,11 +148,11 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 
 	// Try to load a bundle configuration if we're allowed to by the caller (see `./auth_options.go`).
 	if !shouldSkipLoadBundle(cmd.Context()) {
-		err := TryConfigureBundle(cmd, args)
-		if err != nil {
+		b, diags := TryConfigureBundle(cmd)
+		if err := diags.Error(); err != nil {
 			return err
 		}
-		if b := bundle.GetOrNil(cmd.Context()); b != nil {
+		if b != nil {
 			client, err := b.InitializeWorkspaceClient()
 			if err != nil {
 				return err

cmd/root/bundle.go

@@ -4,8 +4,8 @@ import (
 	"context"
 
 	"github.com/databricks/cli/bundle"
-	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/bundle/env"
+	"github.com/databricks/cli/bundle/phases"
 	"github.com/databricks/cli/libs/diag"
 	envlib "github.com/databricks/cli/libs/env"
 	"github.com/spf13/cobra"
@@ -50,87 +50,100 @@ func getProfile(cmd *cobra.Command) (value string) {
 	return envlib.Get(cmd.Context(), "DATABRICKS_CONFIG_PROFILE")
 }
 
-// loadBundle loads the bundle configuration and applies default mutators.
-func loadBundle(cmd *cobra.Command, args []string, load func(ctx context.Context) (*bundle.Bundle, error)) (*bundle.Bundle, error) {
-	ctx := cmd.Context()
-	b, err := load(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	// No bundle is fine in case of `TryConfigureBundle`.
-	if b == nil {
-		return nil, nil
-	}
-
+// configureProfile applies the profile flag to the bundle.
+func configureProfile(cmd *cobra.Command, b *bundle.Bundle) diag.Diagnostics {
 	profile := getProfile(cmd)
-	if profile != "" {
-		diags := bundle.ApplyFunc(ctx, b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-			b.Config.Workspace.Profile = profile
-			return nil
-		})
-		if err := diags.Error(); err != nil {
-			return nil, err
-		}
-	}
-
-	diags := bundle.Apply(ctx, b, bundle.Seq(mutator.DefaultMutators()...))
-	if err := diags.Error(); err != nil {
-		return nil, err
-	}
-
-	return b, nil
-}
-
-// configureBundle loads the bundle configuration and configures it on the command's context.
-func configureBundle(cmd *cobra.Command, args []string, load func(ctx context.Context) (*bundle.Bundle, error)) error {
-	b, err := loadBundle(cmd, args, load)
-	if err != nil {
-		return err
-	}
-
-	// No bundle is fine in case of `TryConfigureBundle`.
-	if b == nil {
+	if profile == "" {
 		return nil
 	}
 
+	return bundle.ApplyFunc(cmd.Context(), b, func(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+		b.Config.Workspace.Profile = profile
+		return nil
+	})
+}
+
+// configureBundle loads the bundle configuration and configures flag values, if any.
+func configureBundle(cmd *cobra.Command, b *bundle.Bundle) (*bundle.Bundle, diag.Diagnostics) {
 	var m bundle.Mutator
-	env := getTarget(cmd)
-	if env == "" {
-		m = mutator.SelectDefaultTarget()
+	if target := getTarget(cmd); target == "" {
+		m = phases.LoadDefaultTarget()
 	} else {
-		m = mutator.SelectTarget(env)
+		m = phases.LoadNamedTarget(target)
 	}
 
+	// Load bundle and select target.
 	ctx := cmd.Context()
 	diags := bundle.Apply(ctx, b, m)
-	if err := diags.Error(); err != nil {
-		return err
+	if diags.HasError() {
+		return nil, diags
 	}
 
-	cmd.SetContext(bundle.Context(ctx, b))
-	return nil
+	// Configure the workspace profile if the flag has been set.
+	diags = diags.Extend(configureProfile(cmd, b))
+	if diags.HasError() {
+		return nil, diags
+	}
+
+	return b, diags
 }
 
 // MustConfigureBundle configures a bundle on the command context.
-func MustConfigureBundle(cmd *cobra.Command, args []string) error {
-	return configureBundle(cmd, args, bundle.MustLoad)
+func MustConfigureBundle(cmd *cobra.Command) (*bundle.Bundle, diag.Diagnostics) {
+	// A bundle may be configured on the context when testing.
+	// If it is, return it immediately.
+	b := bundle.GetOrNil(cmd.Context())
+	if b != nil {
+		return b, nil
+	}
+
+	b, err := bundle.MustLoad(cmd.Context())
+	if err != nil {
+		return nil, diag.FromErr(err)
+	}
+
+	return configureBundle(cmd, b)
 }
 
 // TryConfigureBundle configures a bundle on the command context
 // if there is one, but doesn't fail if there isn't one.
-func TryConfigureBundle(cmd *cobra.Command, args []string) error {
-	return configureBundle(cmd, args, bundle.TryLoad)
+func TryConfigureBundle(cmd *cobra.Command) (*bundle.Bundle, diag.Diagnostics) {
+	// A bundle may be configured on the context when testing.
+	// If it is, return it immediately.
+	b := bundle.GetOrNil(cmd.Context())
+	if b != nil {
+		return b, nil
+	}
+
+	b, err := bundle.TryLoad(cmd.Context())
+	if err != nil {
+		return nil, diag.FromErr(err)
+	}
+
+	// No bundle is fine in this case.
+	if b == nil {
+		return nil, nil
+	}
+
+	return configureBundle(cmd, b)
 }
 
 // targetCompletion executes to autocomplete the argument to the target flag.
 func targetCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-	b, err := loadBundle(cmd, args, bundle.MustLoad)
+	ctx := cmd.Context()
+	b, err := bundle.MustLoad(ctx)
 	if err != nil {
 		cobra.CompErrorln(err.Error())
 		return nil, cobra.ShellCompDirectiveError
 	}
 
+	// Load bundle but don't select a target (we're completing those).
+	diags := bundle.Apply(ctx, b, phases.Load())
+	if err := diags.Error(); err != nil {
+		cobra.CompErrorln(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
 	return maps.Keys(b.Config.Targets), cobra.ShellCompDirectiveDefault
 }
 

cmd/root/bundle_test.go

@@ -2,16 +2,17 @@ package root
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
-	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func setupDatabricksCfg(t *testing.T) {
@@ -37,47 +38,61 @@ func emptyCommand(t *testing.T) *cobra.Command {
 	return cmd
 }
 
-func setup(t *testing.T, cmd *cobra.Command, host string) *bundle.Bundle {
+func setupWithHost(t *testing.T, cmd *cobra.Command, host string) *bundle.Bundle {
 	setupDatabricksCfg(t)
 
 	rootPath := t.TempDir()
-	testutil.Touch(t, rootPath, "databricks.yml")
+	testutil.Chdir(t, rootPath)
 
-	err := configureBundle(cmd, []string{"validate"}, func(_ context.Context) (*bundle.Bundle, error) {
-		return &bundle.Bundle{
-			RootPath: rootPath,
-			Config: config.Root{
-				Bundle: config.Bundle{
-					Name: "test",
-				},
-				Workspace: config.Workspace{
-					Host: host,
-				},
-			},
-		}, nil
-	})
-	assert.NoError(t, err)
-	return bundle.Get(cmd.Context())
+	contents := fmt.Sprintf(`
+workspace:
+  host: %q
+`, host)
+	err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0644)
+	require.NoError(t, err)
+
+	b, diags := MustConfigureBundle(cmd)
+	require.NoError(t, diags.Error())
+	return b
+}
+
+func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) *bundle.Bundle {
+	setupDatabricksCfg(t)
+
+	rootPath := t.TempDir()
+	testutil.Chdir(t, rootPath)
+
+	contents := fmt.Sprintf(`
+workspace:
+  profile: %q
+`, profile)
+	err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0644)
+	require.NoError(t, err)
+
+	b, diags := MustConfigureBundle(cmd)
+	require.NoError(t, diags.Error())
+	return b
 }
 
 func TestBundleConfigureDefault(t *testing.T) {
 	testutil.CleanupEnvironment(t)
 
 	cmd := emptyCommand(t)
-	b := setup(t, cmd, "https://x.com")
-	assert.NotPanics(t, func() {
-		b.WorkspaceClient()
-	})
+	b := setupWithHost(t, cmd, "https://x.com")
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://x.com", client.Config.Host)
 }
 
 func TestBundleConfigureWithMultipleMatches(t *testing.T) {
 	testutil.CleanupEnvironment(t)
 
 	cmd := emptyCommand(t)
-	b := setup(t, cmd, "https://a.com")
-	assert.Panics(t, func() {
-		b.WorkspaceClient()
-	})
+	b := setupWithHost(t, cmd, "https://a.com")
+
+	_, err := b.InitializeWorkspaceClient()
+	assert.ErrorContains(t, err, "multiple profiles matched: PROFILE-1, PROFILE-2")
 }
 
 func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
@@ -85,11 +100,10 @@ func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
 
 	cmd := emptyCommand(t)
 	cmd.Flag("profile").Value.Set("NOEXIST")
+	b := setupWithHost(t, cmd, "https://x.com")
 
-	b := setup(t, cmd, "https://x.com")
-	assert.Panics(t, func() {
-		b.WorkspaceClient()
-	})
+	_, err := b.InitializeWorkspaceClient()
+	assert.ErrorContains(t, err, "has no NOEXIST profile configured")
 }
 
 func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
@@ -97,11 +111,10 @@ func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
 
 	cmd := emptyCommand(t)
 	cmd.Flag("profile").Value.Set("PROFILE-1")
+	b := setupWithHost(t, cmd, "https://x.com")
 
-	b := setup(t, cmd, "https://x.com")
-	assert.PanicsWithError(t, "cannot resolve bundle auth configuration: config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com", func() {
-		b.WorkspaceClient()
-	})
+	_, err := b.InitializeWorkspaceClient()
+	assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
 }
 
 func TestBundleConfigureWithCorrectProfile(t *testing.T) {
@@ -109,35 +122,97 @@ func TestBundleConfigureWithCorrectProfile(t *testing.T) {
 
 	cmd := emptyCommand(t)
 	cmd.Flag("profile").Value.Set("PROFILE-1")
+	b := setupWithHost(t, cmd, "https://a.com")
 
-	b := setup(t, cmd, "https://a.com")
-	assert.NotPanics(t, func() {
-		b.WorkspaceClient()
-	})
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "PROFILE-1", client.Config.Profile)
 }
 
 func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
 	testutil.CleanupEnvironment(t)
-	t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
 
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
 	cmd := emptyCommand(t)
-	b := setup(t, cmd, "https://x.com")
-	assert.PanicsWithError(t, "cannot resolve bundle auth configuration: config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com", func() {
-		b.WorkspaceClient()
-	})
+	b := setupWithHost(t, cmd, "https://x.com")
+
+	_, err := b.InitializeWorkspaceClient()
+	assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
 }
 
 func TestBundleConfigureWithProfileFlagAndEnvVariable(t *testing.T) {
 	testutil.CleanupEnvironment(t)
-	t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
 
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
 	cmd := emptyCommand(t)
 	cmd.Flag("profile").Value.Set("PROFILE-1")
+	b := setupWithHost(t, cmd, "https://a.com")
 
-	b := setup(t, cmd, "https://a.com")
-	assert.NotPanics(t, func() {
-		b.WorkspaceClient()
-	})
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "PROFILE-1", client.Config.Profile)
+}
+
+func TestBundleConfigureProfileDefault(t *testing.T) {
+	testutil.CleanupEnvironment(t)
+
+	// The profile in the databricks.yml file is used
+	cmd := emptyCommand(t)
+	b := setupWithProfile(t, cmd, "PROFILE-1")
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "a", client.Config.Token)
+	assert.Equal(t, "PROFILE-1", client.Config.Profile)
+}
+
+func TestBundleConfigureProfileFlag(t *testing.T) {
+	testutil.CleanupEnvironment(t)
+
+	// The --profile flag takes precedence over the profile in the databricks.yml file
+	cmd := emptyCommand(t)
+	cmd.Flag("profile").Value.Set("PROFILE-2")
+	b := setupWithProfile(t, cmd, "PROFILE-1")
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "PROFILE-2", client.Config.Profile)
+}
+
+func TestBundleConfigureProfileEnvVariable(t *testing.T) {
+	testutil.CleanupEnvironment(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()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "PROFILE-2", client.Config.Profile)
+}
+
+func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
+	testutil.CleanupEnvironment(t)
+
+	// The --profile flag takes precedence over the DATABRICKS_CONFIG_PROFILE environment variable
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
+	cmd := emptyCommand(t)
+	cmd.Flag("profile").Value.Set("PROFILE-2")
+	b := setupWithProfile(t, cmd, "PROFILE-1")
+
+	client, err := b.InitializeWorkspaceClient()
+	require.NoError(t, err)
+	assert.Equal(t, "https://a.com", client.Config.Host)
+	assert.Equal(t, "b", client.Config.Token)
+	assert.Equal(t, "PROFILE-2", client.Config.Profile)
 }
 
 func TestTargetFlagFull(t *testing.T) {
@@ -149,7 +224,7 @@ func TestTargetFlagFull(t *testing.T) {
 	err := cmd.ExecuteContext(ctx)
 	assert.NoError(t, err)
 
-	assert.Equal(t, getTarget(cmd), "development")
+	assert.Equal(t, "development", getTarget(cmd))
 }
 
 func TestTargetFlagShort(t *testing.T) {
@@ -161,7 +236,7 @@ func TestTargetFlagShort(t *testing.T) {
 	err := cmd.ExecuteContext(ctx)
 	assert.NoError(t, err)
 
-	assert.Equal(t, getTarget(cmd), "production")
+	assert.Equal(t, "production", getTarget(cmd))
 }
 
 // TODO: remove when environment flag is fully deprecated
@@ -175,5 +250,5 @@ func TestTargetEnvironmentFlag(t *testing.T) {
 	err := cmd.ExecuteContext(ctx)
 	assert.NoError(t, err)
 
-	assert.Equal(t, getTarget(cmd), "development")
+	assert.Equal(t, "development", getTarget(cmd))
 }