Add more flags to `configure` command (#29)

This commit is contained in:
Kartik Gupta 2022-09-06 16:37:58 +02:00 committed by GitHub
parent 457f3ad3c2
commit 30a7de621a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 104 additions and 43 deletions

View File

@ -15,44 +15,68 @@ import (
)
type Configs struct {
Host string `ini:"host"`
Token string `ini:"token"`
Host string `ini:"host"`
Token string `ini:"token,omitempty"`
Profile string `ini:"-"`
}
var noInteractive bool
var noInteractive, tokenMode bool
func (cfg *Configs) readFromStdin() error {
n, err := fmt.Scanf("%s %s\n", &cfg.Host, &cfg.Token)
func (cfg *Configs) loadNonInteractive(cmd *cobra.Command) error {
host, err := cmd.Flags().GetString("host")
if err != nil || host == "" {
return fmt.Errorf("use --host to specify host in non interactive mode: %w", err)
}
cfg.Host = host
if !tokenMode {
return nil
}
n, err := fmt.Scanf("%s\n", &cfg.Token)
if err != nil {
return err
}
if n != 2 {
return fmt.Errorf("exactly 2 arguments are required")
if n != 1 {
return fmt.Errorf("exactly 1 argument required")
}
return nil
}
func (cfg *Configs) prompt() error {
func (cfg *Configs) loadInteractive(cmd *cobra.Command) error {
res := prompt.Results{}
err := prompt.Questions{prompt.Text{
Key: "host",
Label: "Databricks Host",
Default: func(res prompt.Results) string {
return cfg.Host
},
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
cfg.Host = ans.Value
},
}, prompt.Text{
Key: "token",
Label: "Databricks Token",
Default: func(res prompt.Results) string {
return cfg.Token
},
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
cfg.Token = ans.Value
},
}}.Ask(res)
questions := prompt.Questions{}
host, err := cmd.Flags().GetString("host")
if err != nil || host == "" {
questions = append(questions, prompt.Text{
Key: "host",
Label: "Databricks Host",
Default: func(res prompt.Results) string {
return cfg.Host
},
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
cfg.Host = ans.Value
},
})
} else {
cfg.Host = host
}
if tokenMode {
questions = append(questions, prompt.Text{
Key: "token",
Label: "Databricks Token",
Default: func(res prompt.Results) string {
return cfg.Token
},
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
cfg.Token = ans.Value
},
})
}
err = questions.Ask(res)
if err != nil {
return err
}
@ -70,7 +94,11 @@ var configureCmd = &cobra.Command{
Use: "configure",
Short: "Configure authentication",
RunE: func(cmd *cobra.Command, args []string) error {
var err error
profile, err := cmd.Flags().GetString("profile")
if err != nil {
return fmt.Errorf("read --profile flag: %w", err)
}
path := os.Getenv("DATABRICKS_CONFIG_FILE")
if path == "" {
path, err = os.UserHomeDir()
@ -101,27 +129,30 @@ var configureCmd = &cobra.Command{
if err != nil {
return fmt.Errorf("load config file: %w", err)
}
cfg := &Configs{"", ""}
err = ini_cfg.MapTo(cfg)
cfg := &Configs{"", "", profile}
err = ini_cfg.Section(profile).MapTo(cfg)
if err != nil {
return fmt.Errorf("unmarshal loaded config: %w", err)
}
if noInteractive {
err = cfg.readFromStdin()
err = cfg.loadNonInteractive(cmd)
} else {
err = cfg.prompt()
err = cfg.loadInteractive(cmd)
}
if err != nil {
return fmt.Errorf("reading configs: %w", err)
}
var buffer bytes.Buffer
buffer.WriteString("[DEFAULT]\n")
err = ini_cfg.ReflectFrom(cfg)
err = ini_cfg.Section(profile).ReflectFrom(cfg)
if err != nil {
return fmt.Errorf("marshall config: %w", err)
}
var buffer bytes.Buffer
//The ini library does not write [DEFAULT] header, so we always
//add the [DEFAULT] header explicitly. The section might be empty.
buffer.WriteString("[DEFAULT]\n")
_, err = ini_cfg.WriteTo(&buffer)
if err != nil {
return fmt.Errorf("write config to buffer: %w", err)
@ -137,5 +168,8 @@ var configureCmd = &cobra.Command{
func init() {
root.RootCmd.AddCommand(configureCmd)
configureCmd.Flags().BoolVar(&noInteractive, "no-interactive", false, "Don't show interactive prompts for inputs. Read directly from stdin")
configureCmd.Flags().BoolVarP(&tokenMode, "token", "t", false, "Configure using Databricks Personal Access Token")
configureCmd.Flags().BoolVar(&noInteractive, "no-interactive", false, "Don't show interactive prompts for inputs. Read directly from stdin.")
configureCmd.Flags().String("host", "", "Host to connect to.")
configureCmd.Flags().String("profile", "DEFAULT", "CLI connection profile to use.")
}

View File

@ -45,12 +45,12 @@ func getTempFileWithContent(t *testing.T, tempHomeDir string, content string) *o
func TestDefaultConfigureNoInteractive(t *testing.T) {
ctx := context.Background()
tempHomeDir := setup(t)
inp := getTempFileWithContent(t, tempHomeDir, "host token\n")
inp := getTempFileWithContent(t, tempHomeDir, "token\n")
oldStdin := os.Stdin
defer func() { os.Stdin = oldStdin }()
t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--no-interactive"})
root.RootCmd.SetArgs([]string{"configure", "--token", "--no-interactive", "--host", "host"})
err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err)
@ -76,12 +76,12 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
cfgFileDir := filepath.Join(tempHomeDir, "test")
tests.SetTestEnv(t, "DATABRICKS_CONFIG_FILE", cfgFileDir)
inp := getTempFileWithContent(t, tempHomeDir, "host token\n")
inp := getTempFileWithContent(t, tempHomeDir, "token\n")
oldStdin := os.Stdin
defer func() { os.Stdin = oldStdin }()
t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--no-interactive"})
root.RootCmd.SetArgs([]string{"configure", "--token", "--no-interactive", "--host", "host"})
err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err)
@ -99,3 +99,30 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "token", "token")
}
func TestCustomProfileConfigureNoInteractive(t *testing.T) {
ctx := context.Background()
tempHomeDir := setup(t)
inp := getTempFileWithContent(t, tempHomeDir, "token\n")
oldStdin := os.Stdin
t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--token", "--no-interactive", "--host", "host", "--profile", "CUSTOM"})
err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err)
cfgPath := filepath.Join(tempHomeDir, ".databrickscfg")
_, err = os.Stat(cfgPath)
assert.NoError(t, err)
cfg, err := ini.Load(cfgPath)
assert.NoError(t, err)
defaultSection, err := cfg.GetSection("CUSTOM")
assert.NoError(t, err)
assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "token", "token")
}

2
go.mod
View File

@ -87,5 +87,5 @@ require (
google.golang.org/grpc v1.46.0 // indirect
google.golang.org/protobuf v1.28.0 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)