From f65ab73cc618d33991817ea020e869545a2fb78d Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 8 Jun 2023 22:06:47 +0200 Subject: [PATCH 1/5] Tweak profile prompt This includes the following changes: * Move profile loading code to libs/databrickscfg and add tests * Update prompt label to reflect workspace/account profiles * Start prompt in search mode by default --- cmd/root/auth.go | 112 +++++++++------------------- libs/databrickscfg/profiles.go | 92 +++++++++++++++++++++++ libs/databrickscfg/profiles_test.go | 34 +++++++++ 3 files changed, 160 insertions(+), 78 deletions(-) create mode 100644 libs/databrickscfg/profiles.go create mode 100644 libs/databrickscfg/profiles_test.go diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 39f7bf22ea..56336b1f34 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -5,17 +5,16 @@ import ( "errors" "fmt" "os" - "path/filepath" "strings" "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/manifoldco/promptui" "github.com/spf13/cobra" - "gopkg.in/ini.v1" ) // Placeholders to use as unique keys in context.Context. @@ -41,19 +40,12 @@ func MustAccountClient(cmd *cobra.Command, args []string) error { // 1. only admins will have account configured // 2. 99% of admins will have access to just one account // hence, we don't need to create a special "DEFAULT_ACCOUNT" profile yet - profiles, err := loadProfiles() + _, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchAccountProfiles) if err != nil { return err } - var items []profile - for _, v := range profiles { - if v.AccountID == "" { - continue - } - items = append(items, v) - } - if len(items) == 1 { - cfg.Profile = items[0].Name + if len(profiles) == 1 { + cfg.Profile = profiles[0].Name } } @@ -121,107 +113,71 @@ TRY_AUTH: // or try picking a config profile dynamically return nil } -type profile struct { - Name string - Host string - AccountID string -} +type searchProfiles databrickscfg.Profiles -func (p profile) Cloud() string { - if strings.Contains(p.Host, ".azuredatabricks.net") { - return "Azure" - } - if strings.Contains(p.Host, "gcp.databricks.com") { - return "GCP" - } - return "AWS" -} - -func loadProfiles() (profiles []profile, err error) { - homedir, err := os.UserHomeDir() - if err != nil { - return nil, fmt.Errorf("cannot find homedir: %w", err) - } - file := filepath.Join(homedir, ".databrickscfg") - iniFile, err := ini.Load(file) - if err != nil { - return - } - for _, v := range iniFile.Sections() { - all := v.KeysHash() - host, ok := all["host"] - if !ok { - // invalid profile - continue - } - profiles = append(profiles, profile{ - Name: v.Name(), - Host: host, - AccountID: all["account_id"], - }) - } - return profiles, nil +// Search implements the promptui.Searcher interface. +// This allows the user to immediately starting typing to narrow down the list. +func (s searchProfiles) Search(input string, index int) bool { + input = strings.ToLower(input) + name := strings.ToLower(s[index].Name) + host := strings.ToLower(s[index].Host) + return strings.Contains(name, input) || strings.Contains(host, input) } func askForWorkspaceProfile() (string, error) { - profiles, err := loadProfiles() + file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchWorkspaceProfiles) if err != nil { return "", err } - var items []profile - for _, v := range profiles { - if v.AccountID != "" { - continue - } - items = append(items, v) + // Pick the only profile if there is only one. + if len(profiles) == 1 { + return profiles[0].Name, nil } - label := "~/.databrickscfg profile" i, _, err := (&promptui.Select{ - Label: label, - Items: items, + Label: fmt.Sprintf("Workspace profiles defined in %s", file), + Items: profiles, + Searcher: searchProfiles(profiles).Search, + StartInSearchMode: true, Templates: &promptui.SelectTemplates{ + Label: "{{ . | faint }}", Active: `{{.Name | bold}} ({{.Host|faint}})`, Inactive: `{{.Name}}`, - Selected: fmt.Sprintf(`{{ "%s" | faint }}: {{ .Name | bold }}`, label), + Selected: `{{ "Using workspace profile" | faint }}: {{ .Name | bold }}`, }, Stdin: os.Stdin, }).Run() if err != nil { return "", err } - return items[i].Name, nil + return profiles[i].Name, nil } func askForAccountProfile() (string, error) { - profiles, err := loadProfiles() + file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchAccountProfiles) if err != nil { return "", err } - var items []profile - for _, v := range profiles { - if v.AccountID == "" { - continue - } - items = append(items, v) - } - if len(items) == 1 { - return items[0].Name, nil + // Pick the only profile if there is only one. + if len(profiles) == 1 { + return profiles[0].Name, nil } - label := "~/.databrickscfg profile" i, _, err := (&promptui.Select{ - Label: label, - Items: items, + Label: fmt.Sprintf("Account profiles defined in %s", file), + Items: profiles, + Searcher: searchProfiles(profiles).Search, + StartInSearchMode: true, Templates: &promptui.SelectTemplates{ + Label: "{{ . | faint }}", Active: `{{.Name | bold}} ({{.AccountID|faint}} {{.Cloud|faint}})`, Inactive: `{{.Name}}`, - Selected: fmt.Sprintf(`{{ "%s" | faint }}: {{ .Name | bold }}`, label), + Selected: `{{ "Using account profile" | faint }}: {{ .Name | bold }}`, }, Stdin: os.Stdin, }).Run() if err != nil { return "", err } - return items[i].Name, nil + return profiles[i].Name, nil } func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient { diff --git a/libs/databrickscfg/profiles.go b/libs/databrickscfg/profiles.go new file mode 100644 index 0000000000..16000b3a30 --- /dev/null +++ b/libs/databrickscfg/profiles.go @@ -0,0 +1,92 @@ +package databrickscfg + +import ( + "os" + "strings" + + "github.com/databricks/databricks-sdk-go/config" +) + +// Profile holds a subset of the keys in a databrickscfg profile. +// It should only be used for prompting and filtering. +// Use its name to construct a config.Config. +type Profile struct { + Name string + Host string + AccountID string +} + +func (p Profile) Cloud() string { + cfg := config.Config{Host: p.Host} + switch { + case cfg.IsAws(): + return "AWS" + case cfg.IsAzure(): + return "Azure" + case cfg.IsGcp(): + return "GCP" + default: + return "" + } +} + +type Profiles []Profile + +func (p Profiles) Names() []string { + names := make([]string, len(p)) + for i, v := range p { + names[i] = v.Name + } + return names +} + +type ProfileMatchFunction func(Profile) bool + +func MatchWorkspaceProfiles(p Profile) bool { + return p.AccountID == "" +} + +func MatchAccountProfiles(p Profile) bool { + return p.Host != "" && p.AccountID != "" +} + +const DefaultPath = "~/.databrickscfg" + +func LoadProfiles(path string, fn ProfileMatchFunction) (file string, profiles Profiles, err error) { + f, err := config.LoadFile(path) + if err != nil { + return + } + + homedir, err := os.UserHomeDir() + if err != nil { + return + } + + // Replace homedir with ~ if applicable. + // This is to make the output more readable. + file = f.Path() + if strings.HasPrefix(file, homedir) { + file = "~" + file[len(homedir):] + } + + // Iterate over sections and collect matching profiles. + for _, v := range f.Sections() { + all := v.KeysHash() + host, ok := all["host"] + if !ok { + // invalid profile + continue + } + profile := Profile{ + Name: v.Name(), + Host: host, + AccountID: all["account_id"], + } + if fn(profile) { + profiles = append(profiles, profile) + } + } + + return +} diff --git a/libs/databrickscfg/profiles_test.go b/libs/databrickscfg/profiles_test.go new file mode 100644 index 0000000000..4bc38a8727 --- /dev/null +++ b/libs/databrickscfg/profiles_test.go @@ -0,0 +1,34 @@ +package databrickscfg + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestProfileCloud(t *testing.T) { + assert.Equal(t, Profile{Host: "https://dbc-XXXXXXXX-YYYY.cloud.databricks.com"}.Cloud(), "AWS") + assert.Equal(t, Profile{Host: "https://adb-xxx.y.azuredatabricks.net/"}.Cloud(), "Azure") + assert.Equal(t, Profile{Host: "https://workspace.gcp.databricks.com/"}.Cloud(), "GCP") + assert.Equal(t, Profile{Host: "https://some.invalid.host.com/"}.Cloud(), "AWS") +} + +func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) { + t.Setenv("HOME", "./testdata") + file, _, err := LoadProfiles("./testdata/databrickscfg", func(p Profile) bool { return true }) + require.NoError(t, err) + assert.Equal(t, "~/databrickscfg", file) +} + +func TestLoadProfilesMatchWorkspace(t *testing.T) { + _, profiles, err := LoadProfiles("./testdata/databrickscfg", MatchWorkspaceProfiles) + require.NoError(t, err) + assert.Equal(t, []string{"DEFAULT", "query", "foo1", "foo2"}, profiles.Names()) +} + +func TestLoadProfilesMatchAccount(t *testing.T) { + _, profiles, err := LoadProfiles("./testdata/databrickscfg", MatchAccountProfiles) + require.NoError(t, err) + assert.Equal(t, []string{"acc"}, profiles.Names()) +} From dac98f9efcc430ba2da47f4e0233d54f26ae0688 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 8 Jun 2023 22:13:03 +0200 Subject: [PATCH 2/5] Move search function --- cmd/root/auth.go | 16 ++-------------- libs/databrickscfg/profiles.go | 9 +++++++++ libs/databrickscfg/profiles_test.go | 11 +++++++++++ 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 56336b1f34..b54a758a68 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "os" - "strings" "github.com/databricks/cli/bundle" "github.com/databricks/cli/libs/cmdio" @@ -113,17 +112,6 @@ TRY_AUTH: // or try picking a config profile dynamically return nil } -type searchProfiles databrickscfg.Profiles - -// Search implements the promptui.Searcher interface. -// This allows the user to immediately starting typing to narrow down the list. -func (s searchProfiles) Search(input string, index int) bool { - input = strings.ToLower(input) - name := strings.ToLower(s[index].Name) - host := strings.ToLower(s[index].Host) - return strings.Contains(name, input) || strings.Contains(host, input) -} - func askForWorkspaceProfile() (string, error) { file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchWorkspaceProfiles) if err != nil { @@ -136,7 +124,7 @@ func askForWorkspaceProfile() (string, error) { i, _, err := (&promptui.Select{ Label: fmt.Sprintf("Workspace profiles defined in %s", file), Items: profiles, - Searcher: searchProfiles(profiles).Search, + Searcher: profiles.SearchCaseInsensitive, StartInSearchMode: true, Templates: &promptui.SelectTemplates{ Label: "{{ . | faint }}", @@ -164,7 +152,7 @@ func askForAccountProfile() (string, error) { i, _, err := (&promptui.Select{ Label: fmt.Sprintf("Account profiles defined in %s", file), Items: profiles, - Searcher: searchProfiles(profiles).Search, + Searcher: profiles.SearchCaseInsensitive, StartInSearchMode: true, Templates: &promptui.SelectTemplates{ Label: "{{ . | faint }}", diff --git a/libs/databrickscfg/profiles.go b/libs/databrickscfg/profiles.go index 16000b3a30..60b2a89a2f 100644 --- a/libs/databrickscfg/profiles.go +++ b/libs/databrickscfg/profiles.go @@ -40,6 +40,15 @@ func (p Profiles) Names() []string { return names } +// SearchCaseInsensitive implements the promptui.Searcher interface. +// This allows the user to immediately starting typing to narrow down the list. +func (p Profiles) SearchCaseInsensitive(input string, index int) bool { + input = strings.ToLower(input) + name := strings.ToLower(p[index].Name) + host := strings.ToLower(p[index].Host) + return strings.Contains(name, input) || strings.Contains(host, input) +} + type ProfileMatchFunction func(Profile) bool func MatchWorkspaceProfiles(p Profile) bool { diff --git a/libs/databrickscfg/profiles_test.go b/libs/databrickscfg/profiles_test.go index 4bc38a8727..1032995f63 100644 --- a/libs/databrickscfg/profiles_test.go +++ b/libs/databrickscfg/profiles_test.go @@ -14,6 +14,17 @@ func TestProfileCloud(t *testing.T) { assert.Equal(t, Profile{Host: "https://some.invalid.host.com/"}.Cloud(), "AWS") } +func TestProfilesSearchCaseInsensitive(t *testing.T) { + profiles := Profiles{ + Profile{Name: "foo", Host: "bar"}, + } + assert.True(t, profiles.SearchCaseInsensitive("f", 0)) + assert.True(t, profiles.SearchCaseInsensitive("OO", 0)) + assert.True(t, profiles.SearchCaseInsensitive("b", 0)) + assert.True(t, profiles.SearchCaseInsensitive("AR", 0)) + assert.False(t, profiles.SearchCaseInsensitive("qu", 0)) +} + func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) { t.Setenv("HOME", "./testdata") file, _, err := LoadProfiles("./testdata/databrickscfg", func(p Profile) bool { return true }) From 3729388b7a1e61769763bcc4e3a35b8c5b4b3d20 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 8 Jun 2023 22:47:58 +0200 Subject: [PATCH 3/5] Custom errors; prompt on stderr --- cmd/root/auth.go | 35 +++++++++++++++++++++++++---------- libs/cmdio/io.go | 8 ++++++-- 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index b54a758a68..6787a06bab 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -112,13 +112,23 @@ TRY_AUTH: // or try picking a config profile dynamically return nil } +func transformLoadError(path string, err error) error { + if os.IsNotExist(err) { + return fmt.Errorf("no configuration file found at %s; please create one first", path) + } + return err +} + func askForWorkspaceProfile() (string, error) { - file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchWorkspaceProfiles) + path := databrickscfg.DefaultPath + file, profiles, err := databrickscfg.LoadProfiles(path, databrickscfg.MatchWorkspaceProfiles) if err != nil { - return "", err + return "", transformLoadError(path, err) } - // Pick the only profile if there is only one. - if len(profiles) == 1 { + switch len(profiles) { + case 0: + return "", fmt.Errorf("%s does not contain workspace profiles; please create one first", path) + case 1: return profiles[0].Name, nil } i, _, err := (&promptui.Select{ @@ -132,7 +142,8 @@ func askForWorkspaceProfile() (string, error) { Inactive: `{{.Name}}`, Selected: `{{ "Using workspace profile" | faint }}: {{ .Name | bold }}`, }, - Stdin: os.Stdin, + Stdin: os.Stdin, + Stdout: os.Stderr, }).Run() if err != nil { return "", err @@ -141,12 +152,15 @@ func askForWorkspaceProfile() (string, error) { } func askForAccountProfile() (string, error) { - file, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchAccountProfiles) + path := databrickscfg.DefaultPath + file, profiles, err := databrickscfg.LoadProfiles(path, databrickscfg.MatchAccountProfiles) if err != nil { - return "", err + return "", transformLoadError(path, err) } - // Pick the only profile if there is only one. - if len(profiles) == 1 { + switch len(profiles) { + case 0: + return "", fmt.Errorf("%s does not contain account profiles; please create one first", path) + case 1: return profiles[0].Name, nil } i, _, err := (&promptui.Select{ @@ -160,7 +174,8 @@ func askForAccountProfile() (string, error) { Inactive: `{{.Name}}`, Selected: `{{ "Using account profile" | faint }}: {{ .Name | bold }}`, }, - Stdin: os.Stdin, + Stdin: os.Stdin, + Stdout: os.Stderr, }).Run() if err != nil { return "", err diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index 32637b1d4d..4f2650d5f4 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -10,7 +10,6 @@ import ( "github.com/briandowns/spinner" "github.com/databricks/cli/libs/flags" - "github.com/fatih/color" "github.com/manifoldco/promptui" "github.com/mattn/go-isatty" "golang.org/x/exp/slices" @@ -31,8 +30,13 @@ type cmdIO struct { } func NewIO(outputFormat flags.Output, in io.Reader, out io.Writer, err io.Writer, template string) *cmdIO { + // The check below is similar to color.NoColor but uses the specified err writer. + dumb := os.Getenv("NO_COLOR") != "" || os.Getenv("TERM") == "dumb" + if f, ok := err.(*os.File); ok && !dumb { + dumb = !isatty.IsTerminal(f.Fd()) && !isatty.IsCygwinTerminal(f.Fd()) + } return &cmdIO{ - interactive: !color.NoColor, + interactive: !dumb, outputFormat: outputFormat, template: template, in: in, From 9e0fbe069071cc50ac1799e6e74f0ca47fec88a3 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 8 Jun 2023 22:50:39 +0200 Subject: [PATCH 4/5] Fix test on Windows --- libs/databrickscfg/profiles_test.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/databrickscfg/profiles_test.go b/libs/databrickscfg/profiles_test.go index 1032995f63..582c6658e7 100644 --- a/libs/databrickscfg/profiles_test.go +++ b/libs/databrickscfg/profiles_test.go @@ -1,6 +1,7 @@ package databrickscfg import ( + "runtime" "testing" "github.com/stretchr/testify/assert" @@ -26,7 +27,11 @@ func TestProfilesSearchCaseInsensitive(t *testing.T) { } func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) { - t.Setenv("HOME", "./testdata") + if runtime.GOOS == "windows" { + t.Setenv("USERPROFILE", "./testdata") + } else { + t.Setenv("HOME", "./testdata") + } file, _, err := LoadProfiles("./testdata/databrickscfg", func(p Profile) bool { return true }) require.NoError(t, err) assert.Equal(t, "~/databrickscfg", file) From 5b3107cbec9813a011660ab8bd8bed36f2462baf Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 8 Jun 2023 23:02:06 +0200 Subject: [PATCH 5/5] Line length --- cmd/root/auth.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 6787a06bab..61068ab384 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -39,7 +39,10 @@ func MustAccountClient(cmd *cobra.Command, args []string) error { // 1. only admins will have account configured // 2. 99% of admins will have access to just one account // hence, we don't need to create a special "DEFAULT_ACCOUNT" profile yet - _, profiles, err := databrickscfg.LoadProfiles(databrickscfg.DefaultPath, databrickscfg.MatchAccountProfiles) + _, profiles, err := databrickscfg.LoadProfiles( + databrickscfg.DefaultPath, + databrickscfg.MatchAccountProfiles, + ) if err != nil { return err }