From eac176ddac3db6b3f891ea1f9382c1824df6631d Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Tue, 11 Jul 2023 16:01:00 +0200 Subject: [PATCH 1/5] Correctly use --profile flag passed for all bundle commands --- bundle/config/workspace.go | 12 ++++++++---- cmd/root/bundle.go | 19 +++++++++++++++++++ libs/databrickscfg/loader.go | 2 +- libs/databrickscfg/ops.go | 24 ++++++++++++++++++++++++ 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 8a1205b3b..08e30ce65 100644 --- a/bundle/config/workspace.go +++ b/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 != "" { + err := databrickscfg.ValidateConfigAndProfileHost(&cfg, w.Profile) + if err != nil { + return nil, err + } + } + return databricks.NewWorkspaceClient(&cfg) } diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go index 737651242..8eab7c2c7 100644 --- a/cmd/root/bundle.go +++ b/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 { diff --git a/libs/databrickscfg/loader.go b/libs/databrickscfg/loader.go index 8179703a3..05698eb48 100644 --- a/libs/databrickscfg/loader.go +++ b/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 { diff --git a/libs/databrickscfg/ops.go b/libs/databrickscfg/ops.go index 4a4a27b06..59e95d641 100644 --- a/libs/databrickscfg/ops.go +++ b/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 has %s host 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. From 3b02460ec3ae8842a242901d0169a52f8e4a77cf Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Tue, 11 Jul 2023 17:18:29 +0200 Subject: [PATCH 2/5] Update libs/databrickscfg/ops.go Co-authored-by: Pieter Noordhuis --- libs/databrickscfg/ops.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/databrickscfg/ops.go b/libs/databrickscfg/ops.go index 59e95d641..c2d6e9fa1 100644 --- a/libs/databrickscfg/ops.go +++ b/libs/databrickscfg/ops.go @@ -147,7 +147,7 @@ func ValidateConfigAndProfileHost(cfg *databricks.Config, profile string) error hostFromProfile := normalizeHost(match.Key("host").Value()) if hostFromProfile != "" && host != "" && hostFromProfile != host { - return fmt.Errorf("config host mismatch, profile has %s host but CLI configured to use %s ", hostFromProfile, host) + return fmt.Errorf("config host mismatch: profile uses host %s, but CLI configured to use %s", hostFromProfile, host) } return nil From d416c9c37c773089738ace759c6d3890b71eea07 Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Tue, 11 Jul 2023 17:19:48 +0200 Subject: [PATCH 3/5] check cfg.Host for empty --- bundle/config/workspace.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 08e30ce65..4b09237da 100644 --- a/bundle/config/workspace.go +++ b/bundle/config/workspace.go @@ -88,7 +88,7 @@ func (w *Workspace) Client() (*databricks.WorkspaceClient, error) { } } - if w.Profile != "" { + if w.Profile != "" && cfg.Host != "" { err := databrickscfg.ValidateConfigAndProfileHost(&cfg, w.Profile) if err != nil { return nil, err From fa2bf0afbc5a40c0f7608b8a2acb02e6130a1ade Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Wed, 12 Jul 2023 13:08:03 +0200 Subject: [PATCH 4/5] Added unit tests --- cmd/root/bundle_test.go | 119 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100644 cmd/root/bundle_test.go diff --git a/cmd/root/bundle_test.go b/cmd/root/bundle_test.go new file mode 100644 index 000000000..8dc771bd4 --- /dev/null +++ b/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() + }) +} From a5e49f02b6cbb69e976008b45a3ef05dca94f0d1 Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Wed, 12 Jul 2023 13:10:05 +0200 Subject: [PATCH 5/5] typo fix --- bundle/config/workspace.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 4b09237da..1b6dc4cd5 100644 --- a/bundle/config/workspace.go +++ b/bundle/config/workspace.go @@ -88,7 +88,7 @@ func (w *Workspace) Client() (*databricks.WorkspaceClient, error) { } } - if w.Profile != "" && cfg.Host != "" { + if w.Profile != "" && w.Host != "" { err := databrickscfg.ValidateConfigAndProfileHost(&cfg, w.Profile) if err != nil { return nil, err