Correctly use --profile flag passed for all bundle commands (#571)

## Changes
Correctly use --profile flag passed for all bundle commands.

Also adds a validation that if bundle configured host mismatches
provided profile, it throws an error.

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

bundle/config/workspace.go

@@ -75,12 +75,9 @@ func (w *Workspace) Client() (*databricks.WorkspaceClient, error) {
 		AzureLoginAppID:  w.AzureLoginAppID,
 	}
 
-	// HACKY fix to not used host based auth when the profile is already set
-	profile := os.Getenv("DATABRICKS_CONFIG_PROFILE")
-
 	// If only the host is configured, we try and unambiguously match it to
 	// a profile in the user's databrickscfg file. Override the default loaders.
-	if w.Host != "" && w.Profile == "" && profile == "" {
+	if w.Host != "" && w.Profile == "" {
 		cfg.Loaders = []config.Loader{
 			// Load auth creds from env vars
 			config.ConfigAttributes,
@@ -91,6 +88,13 @@ func (w *Workspace) Client() (*databricks.WorkspaceClient, error) {
 		}
 	}
 
+	if w.Profile != "" && w.Host != "" {
+		err := databrickscfg.ValidateConfigAndProfileHost(&cfg, w.Profile)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return databricks.NewWorkspaceClient(&cfg)
 }
 

cmd/root/bundle.go

@@ -26,6 +26,20 @@ func getEnvironment(cmd *cobra.Command) (value string) {
 	return os.Getenv(envName)
 }
 
+func getProfile(cmd *cobra.Command) (value string) {
+	// The command line flag takes precedence.
+	flag := cmd.Flag("profile")
+	if flag != nil {
+		value = flag.Value.String()
+		if value != "" {
+			return
+		}
+	}
+
+	// If it's not set, use the environment variable.
+	return os.Getenv("DATABRICKS_CONFIG_PROFILE")
+}
+
 // loadBundle loads the bundle configuration and applies default mutators.
 func loadBundle(cmd *cobra.Command, args []string, load func() (*bundle.Bundle, error)) (*bundle.Bundle, error) {
 	b, err := load()
@@ -38,6 +52,11 @@ func loadBundle(cmd *cobra.Command, args []string, load func() (*bundle.Bundle,
 		return nil, nil
 	}
 
+	profile := getProfile(cmd)
+	if profile != "" {
+		b.Config.Workspace.Profile = profile
+	}
+
 	ctx := cmd.Context()
 	err = bundle.Apply(ctx, b, bundle.Seq(mutator.DefaultMutators()...))
 	if err != nil {

cmd/root/bundle_test.go

@@ -0,0 +1,119 @@
+package root
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/stretchr/testify/assert"
+)
+
+func setupDatabricksCfg(t *testing.T) {
+	tempHomeDir := t.TempDir()
+	homeEnvVar := "HOME"
+	if runtime.GOOS == "windows" {
+		homeEnvVar = "USERPROFILE"
+	}
+
+	cfg := []byte("[PROFILE-1]\nhost = https://a.com\ntoken = a\n[PROFILE-2]\nhost = https://a.com\ntoken = b\n")
+	err := os.WriteFile(filepath.Join(tempHomeDir, ".databrickscfg"), cfg, 0644)
+	assert.NoError(t, err)
+
+	t.Setenv("DATABRICKS_CONFIG_FILE", "")
+	t.Setenv(homeEnvVar, tempHomeDir)
+}
+
+func setup(t *testing.T, host string) *bundle.Bundle {
+	setupDatabricksCfg(t)
+
+	ctx := context.Background()
+	RootCmd.SetContext(ctx)
+	_, err := initializeLogger(ctx)
+	assert.NoError(t, err)
+
+	err = configureBundle(RootCmd, []string{"validate"}, func() (*bundle.Bundle, error) {
+		return &bundle.Bundle{
+			Config: config.Root{
+				Bundle: config.Bundle{
+					Name: "test",
+				},
+				Workspace: config.Workspace{
+					Host: host,
+				},
+			},
+		}, nil
+	})
+	assert.NoError(t, err)
+
+	return bundle.Get(RootCmd.Context())
+}
+
+func TestBundleConfigureDefault(t *testing.T) {
+	b := setup(t, "https://x.com")
+	assert.NotPanics(t, func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithMultipleMatches(t *testing.T) {
+	b := setup(t, "https://a.com")
+	assert.Panics(t, func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
+	RootCmd.Flag("profile").Value.Set("NOEXIST")
+
+	b := setup(t, "https://x.com")
+	assert.PanicsWithError(t, "no matching config profiles found", func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
+	RootCmd.Flag("profile").Value.Set("PROFILE-1")
+
+	b := setup(t, "https://x.com")
+	assert.PanicsWithError(t, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com", func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithCorrectProfile(t *testing.T) {
+	RootCmd.Flag("profile").Value.Set("PROFILE-1")
+
+	b := setup(t, "https://a.com")
+	assert.NotPanics(t, func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
+	t.Cleanup(func() {
+		t.Setenv("DATABRICKS_CONFIG_PROFILE", "")
+	})
+
+	b := setup(t, "https://x.com")
+	assert.PanicsWithError(t, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com", func() {
+		b.WorkspaceClient()
+	})
+}
+
+func TestBundleConfigureWithProfileFlagAndEnvVariable(t *testing.T) {
+	t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
+	t.Cleanup(func() {
+		t.Setenv("DATABRICKS_CONFIG_PROFILE", "")
+	})
+	RootCmd.Flag("profile").Value.Set("PROFILE-1")
+
+	b := setup(t, "https://a.com")
+	assert.NotPanics(t, func() {
+		b.WorkspaceClient()
+	})
+}

libs/databrickscfg/loader.go

@@ -90,7 +90,7 @@ func (l profileFromHostLoader) Configure(cfg *config.Config) error {
 	}
 	if err, ok := err.(errMultipleProfiles); ok {
 		return fmt.Errorf(
-			"%s: %w: please set DATABRICKS_CONFIG_PROFILE to specify one",
+			"%s: %w: please set DATABRICKS_CONFIG_PROFILE or provide --profile flag to specify one",
 			host, err)
 	}
 	if err != nil {

libs/databrickscfg/ops.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"gopkg.in/ini.v1"
 )
@@ -129,6 +130,29 @@ func SaveToProfile(ctx context.Context, cfg *config.Config) error {
 	return configFile.SaveTo(configFile.Path())
 }
 
+func ValidateConfigAndProfileHost(cfg *databricks.Config, profile string) error {
+	configFile, err := config.LoadFile(cfg.ConfigFile)
+	if err != nil {
+		return fmt.Errorf("cannot parse config file: %w", err)
+	}
+	// Normalized version of the configured host.
+	host := normalizeHost(cfg.Host)
+	match, err := findMatchingProfile(configFile, func(s *ini.Section) bool {
+		return profile == s.Name()
+	})
+
+	if err != nil {
+		return err
+	}
+
+	hostFromProfile := normalizeHost(match.Key("host").Value())
+	if hostFromProfile != "" && host != "" && hostFromProfile != host {
+		return fmt.Errorf("config host mismatch: profile uses host %s, but CLI configured to use %s", hostFromProfile, host)
+	}
+
+	return nil
+}
+
 func init() {
 	// We document databrickscfg files with a [DEFAULT] header and wish to keep it that way.
 	// This, however, does mean we emit a [DEFAULT] section even if it's empty.