From 30a7de621aa2d56b7389ec7392254f4c34bc1660 Mon Sep 17 00:00:00 2001 From: Kartik Gupta <88345179+kartikgupta-db@users.noreply.github.com> Date: Tue, 6 Sep 2022 16:37:58 +0200 Subject: [PATCH] Add more flags to `configure` command (#29) --- cmd/configure/configure.go | 106 +++++++++++++++++++++----------- cmd/configure/configure_test.go | 39 ++++++++++-- go.mod | 2 +- 3 files changed, 104 insertions(+), 43 deletions(-) diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index e58e5e14..1d14cbe4 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -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.") } diff --git a/cmd/configure/configure_test.go b/cmd/configure/configure_test.go index 1c335c8a..c8ae08f0 100644 --- a/cmd/configure/configure_test.go +++ b/cmd/configure/configure_test.go @@ -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") +} diff --git a/go.mod b/go.mod index 8444b8e0..22c51dd8 100644 --- a/go.mod +++ b/go.mod @@ -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 )