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 { type Configs struct {
Host string `ini:"host"` Host string `ini:"host"`
Token string `ini:"token"` Token string `ini:"token,omitempty"`
Profile string `ini:"-"`
} }
var noInteractive bool var noInteractive, tokenMode bool
func (cfg *Configs) readFromStdin() error { func (cfg *Configs) loadNonInteractive(cmd *cobra.Command) error {
n, err := fmt.Scanf("%s %s\n", &cfg.Host, &cfg.Token) 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 { if err != nil {
return err return err
} }
if n != 2 { if n != 1 {
return fmt.Errorf("exactly 2 arguments are required") return fmt.Errorf("exactly 1 argument required")
} }
return nil return nil
} }
func (cfg *Configs) prompt() error { func (cfg *Configs) loadInteractive(cmd *cobra.Command) error {
res := prompt.Results{} res := prompt.Results{}
err := prompt.Questions{prompt.Text{ questions := prompt.Questions{}
Key: "host",
Label: "Databricks Host", host, err := cmd.Flags().GetString("host")
Default: func(res prompt.Results) string { if err != nil || host == "" {
return cfg.Host questions = append(questions, prompt.Text{
}, Key: "host",
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) { Label: "Databricks Host",
cfg.Host = ans.Value Default: func(res prompt.Results) string {
}, return cfg.Host
}, prompt.Text{ },
Key: "token", Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
Label: "Databricks Token", cfg.Host = ans.Value
Default: func(res prompt.Results) string { },
return cfg.Token })
}, } else {
Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) { cfg.Host = host
cfg.Token = ans.Value }
},
}}.Ask(res) 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 { if err != nil {
return err return err
} }
@ -70,7 +94,11 @@ var configureCmd = &cobra.Command{
Use: "configure", Use: "configure",
Short: "Configure authentication", Short: "Configure authentication",
RunE: func(cmd *cobra.Command, args []string) error { 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") path := os.Getenv("DATABRICKS_CONFIG_FILE")
if path == "" { if path == "" {
path, err = os.UserHomeDir() path, err = os.UserHomeDir()
@ -101,27 +129,30 @@ var configureCmd = &cobra.Command{
if err != nil { if err != nil {
return fmt.Errorf("load config file: %w", err) return fmt.Errorf("load config file: %w", err)
} }
cfg := &Configs{"", ""} cfg := &Configs{"", "", profile}
err = ini_cfg.MapTo(cfg) err = ini_cfg.Section(profile).MapTo(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unmarshal loaded config: %w", err) return fmt.Errorf("unmarshal loaded config: %w", err)
} }
if noInteractive { if noInteractive {
err = cfg.readFromStdin() err = cfg.loadNonInteractive(cmd)
} else { } else {
err = cfg.prompt() err = cfg.loadInteractive(cmd)
} }
if err != nil { if err != nil {
return fmt.Errorf("reading configs: %w", err) return fmt.Errorf("reading configs: %w", err)
} }
var buffer bytes.Buffer err = ini_cfg.Section(profile).ReflectFrom(cfg)
buffer.WriteString("[DEFAULT]\n")
err = ini_cfg.ReflectFrom(cfg)
if err != nil { if err != nil {
return fmt.Errorf("marshall config: %w", err) 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) _, err = ini_cfg.WriteTo(&buffer)
if err != nil { if err != nil {
return fmt.Errorf("write config to buffer: %w", err) return fmt.Errorf("write config to buffer: %w", err)
@ -137,5 +168,8 @@ var configureCmd = &cobra.Command{
func init() { func init() {
root.RootCmd.AddCommand(configureCmd) 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) { func TestDefaultConfigureNoInteractive(t *testing.T) {
ctx := context.Background() ctx := context.Background()
tempHomeDir := setup(t) tempHomeDir := setup(t)
inp := getTempFileWithContent(t, tempHomeDir, "host token\n") inp := getTempFileWithContent(t, tempHomeDir, "token\n")
oldStdin := os.Stdin oldStdin := os.Stdin
defer func() { os.Stdin = oldStdin }() t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp 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) err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err) assert.NoError(t, err)
@ -76,12 +76,12 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
cfgFileDir := filepath.Join(tempHomeDir, "test") cfgFileDir := filepath.Join(tempHomeDir, "test")
tests.SetTestEnv(t, "DATABRICKS_CONFIG_FILE", cfgFileDir) tests.SetTestEnv(t, "DATABRICKS_CONFIG_FILE", cfgFileDir)
inp := getTempFileWithContent(t, tempHomeDir, "host token\n") inp := getTempFileWithContent(t, tempHomeDir, "token\n")
oldStdin := os.Stdin oldStdin := os.Stdin
defer func() { os.Stdin = oldStdin }() t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp 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) err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err) assert.NoError(t, err)
@ -99,3 +99,30 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
assertKeyValueInSection(t, defaultSection, "host", "host") assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "token", "token") 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/grpc v1.46.0 // indirect
google.golang.org/protobuf v1.28.0 // indirect google.golang.org/protobuf v1.28.0 // indirect
gopkg.in/yaml.v2 v2.4.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
) )