Tweak profile prompt (#454)

## Changes

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
* Custom error if `~/.databrickscfg` doesn't exist
* Custom error if `~/.databrickscfg` doesn't contain profiles
* Use stderr for prompt so that stdout redirection works (e.g. with `jq` or `jless`)

## Tests

* New unit tests pass
* Manual tests for both workspace and account commands
* Search-by-default is really nice if you have many profiles
This commit is contained in:
Pieter Noordhuis 2023-06-09 13:56:35 +02:00 committed by GitHub
parent 894d25e434
commit e4415bfbcf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 202 additions and 85 deletions

View File

@ -5,17 +5,15 @@ import (
"errors" "errors"
"fmt" "fmt"
"os" "os"
"path/filepath"
"strings"
"github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle"
"github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/databrickscfg"
"github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config"
"github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/iam"
"github.com/manifoldco/promptui" "github.com/manifoldco/promptui"
"github.com/spf13/cobra" "github.com/spf13/cobra"
"gopkg.in/ini.v1"
) )
// Placeholders to use as unique keys in context.Context. // Placeholders to use as unique keys in context.Context.
@ -41,19 +39,15 @@ func MustAccountClient(cmd *cobra.Command, args []string) error {
// 1. only admins will have account configured // 1. only admins will have account configured
// 2. 99% of admins will have access to just one account // 2. 99% of admins will have access to just one account
// hence, we don't need to create a special "DEFAULT_ACCOUNT" profile yet // 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 { if err != nil {
return err return err
} }
var items []profile if len(profiles) == 1 {
for _, v := range profiles { cfg.Profile = profiles[0].Name
if v.AccountID == "" {
continue
}
items = append(items, v)
}
if len(items) == 1 {
cfg.Profile = items[0].Name
} }
} }
@ -121,107 +115,75 @@ TRY_AUTH: // or try picking a config profile dynamically
return nil return nil
} }
type profile struct { func transformLoadError(path string, err error) error {
Name string if os.IsNotExist(err) {
Host string return fmt.Errorf("no configuration file found at %s; please create one first", path)
AccountID string
} }
return err
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
} }
func askForWorkspaceProfile() (string, error) { func askForWorkspaceProfile() (string, error) {
profiles, err := loadProfiles() path := databrickscfg.DefaultPath
file, profiles, err := databrickscfg.LoadProfiles(path, databrickscfg.MatchWorkspaceProfiles)
if err != nil { if err != nil {
return "", err return "", transformLoadError(path, err)
} }
var items []profile switch len(profiles) {
for _, v := range profiles { case 0:
if v.AccountID != "" { return "", fmt.Errorf("%s does not contain workspace profiles; please create one first", path)
continue case 1:
return profiles[0].Name, nil
} }
items = append(items, v)
}
label := "~/.databrickscfg profile"
i, _, err := (&promptui.Select{ i, _, err := (&promptui.Select{
Label: label, Label: fmt.Sprintf("Workspace profiles defined in %s", file),
Items: items, Items: profiles,
Searcher: profiles.SearchCaseInsensitive,
StartInSearchMode: true,
Templates: &promptui.SelectTemplates{ Templates: &promptui.SelectTemplates{
Label: "{{ . | faint }}",
Active: `{{.Name | bold}} ({{.Host|faint}})`, Active: `{{.Name | bold}} ({{.Host|faint}})`,
Inactive: `{{.Name}}`, Inactive: `{{.Name}}`,
Selected: fmt.Sprintf(`{{ "%s" | faint }}: {{ .Name | bold }}`, label), Selected: `{{ "Using workspace profile" | faint }}: {{ .Name | bold }}`,
}, },
Stdin: os.Stdin, Stdin: os.Stdin,
Stdout: os.Stderr,
}).Run() }).Run()
if err != nil { if err != nil {
return "", err return "", err
} }
return items[i].Name, nil return profiles[i].Name, nil
} }
func askForAccountProfile() (string, error) { func askForAccountProfile() (string, error) {
profiles, err := loadProfiles() path := databrickscfg.DefaultPath
file, profiles, err := databrickscfg.LoadProfiles(path, databrickscfg.MatchAccountProfiles)
if err != nil { if err != nil {
return "", err return "", transformLoadError(path, err)
} }
var items []profile switch len(profiles) {
for _, v := range profiles { case 0:
if v.AccountID == "" { return "", fmt.Errorf("%s does not contain account profiles; please create one first", path)
continue case 1:
return profiles[0].Name, nil
} }
items = append(items, v)
}
if len(items) == 1 {
return items[0].Name, nil
}
label := "~/.databrickscfg profile"
i, _, err := (&promptui.Select{ i, _, err := (&promptui.Select{
Label: label, Label: fmt.Sprintf("Account profiles defined in %s", file),
Items: items, Items: profiles,
Searcher: profiles.SearchCaseInsensitive,
StartInSearchMode: true,
Templates: &promptui.SelectTemplates{ Templates: &promptui.SelectTemplates{
Label: "{{ . | faint }}",
Active: `{{.Name | bold}} ({{.AccountID|faint}} {{.Cloud|faint}})`, Active: `{{.Name | bold}} ({{.AccountID|faint}} {{.Cloud|faint}})`,
Inactive: `{{.Name}}`, Inactive: `{{.Name}}`,
Selected: fmt.Sprintf(`{{ "%s" | faint }}: {{ .Name | bold }}`, label), Selected: `{{ "Using account profile" | faint }}: {{ .Name | bold }}`,
}, },
Stdin: os.Stdin, Stdin: os.Stdin,
Stdout: os.Stderr,
}).Run() }).Run()
if err != nil { if err != nil {
return "", err return "", err
} }
return items[i].Name, nil return profiles[i].Name, nil
} }
func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient { func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient {

View File

@ -10,7 +10,6 @@ import (
"github.com/briandowns/spinner" "github.com/briandowns/spinner"
"github.com/databricks/cli/libs/flags" "github.com/databricks/cli/libs/flags"
"github.com/fatih/color"
"github.com/manifoldco/promptui" "github.com/manifoldco/promptui"
"github.com/mattn/go-isatty" "github.com/mattn/go-isatty"
"golang.org/x/exp/slices" "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 { 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{ return &cmdIO{
interactive: !color.NoColor, interactive: !dumb,
outputFormat: outputFormat, outputFormat: outputFormat,
template: template, template: template,
in: in, in: in,

View File

@ -0,0 +1,101 @@
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
}
// 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 {
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
}

View File

@ -0,0 +1,50 @@
package databrickscfg
import (
"runtime"
"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 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) {
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)
}
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())
}