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

cmd/configure/configure.go

@@ -16,25 +16,40 @@ import (
 
 type Configs struct {
 	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 {
 		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
 }
 
-func (cfg *Configs) prompt() error {
+func (cfg *Configs) loadInteractive(cmd *cobra.Command) error {
 	res := prompt.Results{}
-	err := prompt.Questions{prompt.Text{
+	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 {
@@ -43,7 +58,13 @@ func (cfg *Configs) prompt() error {
 			Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
 				cfg.Host = ans.Value
 			},
-	}, prompt.Text{
+		})
+	} else {
+		cfg.Host = host
+	}
+
+	if tokenMode {
+		questions = append(questions, prompt.Text{
 			Key:   "token",
 			Label: "Databricks Token",
 			Default: func(res prompt.Results) string {
@@ -52,7 +73,10 @@ func (cfg *Configs) prompt() error {
 			Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
 				cfg.Token = ans.Value
 			},
-	}}.Ask(res)
+		})
+	}
+
+	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{"", ""}
+		cfg := &Configs{"", "", profile}
-		err = ini_cfg.MapTo(cfg)
+		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
+		err = ini_cfg.Section(profile).ReflectFrom(cfg)
-		buffer.WriteString("[DEFAULT]\n")
-		err = ini_cfg.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.")
 }

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")
+}

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
 )