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>
This commit is contained in:
Andrew Nester 2023-07-12 14:09:25 +02:00 committed by GitHub
parent 14cfc80666
commit 650fb0e8b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 171 additions and 5 deletions

View File

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

View File

@ -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 {

119
cmd/root/bundle_test.go Normal file
View File

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

View File

@ -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 {

View File

@ -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.