Add `--configure-cluster` flag to configure command (#1005)

## Changes

This breaks out the flags into a separate struct to make it easier to
pass around.

If specified, the flag calls into the `cfgpicker` to prompt for a
cluster to associated with the profile.

## Tests

Existing tests pass; added one for host validation.

---------

Co-authored-by: Miles Yucht <miles@databricks.com>

cmd/configure/configure.go

@@ -1,70 +1,23 @@
 package configure
 
 import (
-	"context"
 	"fmt"
-	"net/url"
 
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg"
+	"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/spf13/cobra"
 )
 
-func validateHost(s string) error {
-	u, err := url.Parse(s)
-	if err != nil {
-		return err
-	}
-	if u.Host == "" || u.Scheme != "https" {
-		return fmt.Errorf("must start with https://")
-	}
-	if u.Path != "" && u.Path != "/" {
-		return fmt.Errorf("must use empty path")
-	}
-	return nil
-}
-
-func configureFromFlags(cmd *cobra.Command, ctx context.Context, cfg *config.Config) error {
-	// Configure profile name if set.
-	profile, err := cmd.Flags().GetString("profile")
-	if err != nil {
-		return fmt.Errorf("read --profile flag: %w", err)
-	}
-	if profile != "" {
-		cfg.Profile = profile
-	}
-
-	// Configure host if set.
-	host, err := cmd.Flags().GetString("host")
-	if err != nil {
-		return fmt.Errorf("read --host flag: %w", err)
-	}
-	if host != "" {
-		cfg.Host = host
-	}
-
-	// Validate host if set.
-	if cfg.Host != "" {
-		err = validateHost(cfg.Host)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func configureInteractive(cmd *cobra.Command, ctx context.Context, cfg *config.Config) error {
-	err := configureFromFlags(cmd, ctx, cfg)
-	if err != nil {
-		return err
-	}
+func configureInteractive(cmd *cobra.Command, flags *configureFlags, cfg *config.Config) error {
+	ctx := cmd.Context()
 
 	// Ask user to specify the host if not already set.
 	if cfg.Host == "" {
 		prompt := cmdio.Prompt(ctx)
-		prompt.Label = "Databricks Host"
+		prompt.Label = "Databricks host"
 		prompt.Default = "https://"
 		prompt.AllowEdit = true
 		prompt.Validate = validateHost
@@ -78,7 +31,7 @@ func configureInteractive(cmd *cobra.Command, ctx context.Context, cfg *config.C
 	// Ask user to specify the token is not already set.
 	if cfg.Token == "" {
 		prompt := cmdio.Prompt(ctx)
-		prompt.Label = "Personal Access Token"
+		prompt.Label = "Personal access token"
 		prompt.Mask = '*'
 		out, err := prompt.Run()
 		if err != nil {
@@ -87,19 +40,32 @@ func configureInteractive(cmd *cobra.Command, ctx context.Context, cfg *config.C
 		cfg.Token = out
 	}
 
+	// Ask user to specify a cluster if not already set.
+	if flags.ConfigureCluster && cfg.ClusterID == "" {
+		w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg))
+		if err != nil {
+			return err
+		}
+		clusterID, err := cfgpickers.AskForCluster(cmd.Context(), w)
+		if err != nil {
+			return err
+		}
+		cfg.ClusterID = clusterID
+	}
+
 	return nil
 }
 
-func configureNonInteractive(cmd *cobra.Command, ctx context.Context, cfg *config.Config) error {
-	err := configureFromFlags(cmd, ctx, cfg)
-	if err != nil {
-		return err
-	}
-
+func configureNonInteractive(cmd *cobra.Command, flags *configureFlags, cfg *config.Config) error {
 	if cfg.Host == "" {
 		return fmt.Errorf("host must be set in non-interactive mode")
 	}
 
+	// Check presence of cluster ID before reading token to fail fast.
+	if flags.ConfigureCluster && cfg.ClusterID == "" {
+		return fmt.Errorf("cluster ID must be set in non-interactive mode")
+	}
+
 	// Read token from stdin if not already set.
 	if cfg.Token == "" {
 		_, err := fmt.Fscanf(cmd.InOrStdin(), "%s\n", &cfg.Token)
@@ -117,21 +83,16 @@ func newConfigureCommand() *cobra.Command {
 		Short: "Configure authentication",
 		Long: `Configure authentication.
 
-	This command adds a profile to your ~/.databrickscfg file.
-	You can write to a different file by setting the DATABRICKS_CONFIG_FILE environment variable.
+This command adds a profile to your ~/.databrickscfg file.
+You can write to a different file by setting the DATABRICKS_CONFIG_FILE environment variable.
 
-	If this command is invoked in non-interactive mode, it will read the token from stdin.
-	The host must be specified with the --host flag.
+If this command is invoked in non-interactive mode, it will read the token from stdin.
+The host must be specified with the --host flag or the DATABRICKS_HOST environment variable.
 		`,
 	}
 
-	cmd.Flags().String("host", "", "Databricks workspace host.")
-	cmd.Flags().String("profile", "DEFAULT", "Name for the connection profile to configure.")
-
-	// Include token flag for compatibility with the legacy CLI.
-	// It doesn't actually do anything because we always use PATs.
-	cmd.Flags().Bool("token", true, "Configure using Databricks Personal Access Token")
-	cmd.Flags().MarkHidden("token")
+	var flags configureFlags
+	flags.Register(cmd)
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		var cfg config.Config
@@ -142,15 +103,28 @@ func newConfigureCommand() *cobra.Command {
 			return fmt.Errorf("unable to instantiate configuration from environment variables: %w", err)
 		}
 
-		ctx := cmd.Context()
-		interactive := cmdio.IsInTTY(ctx) && cmdio.IsOutTTY(ctx)
-		var fn func(*cobra.Command, context.Context, *config.Config) error
-		if interactive {
-			fn = configureInteractive
-		} else {
-			fn = configureNonInteractive
+		// Populate configuration from flags (if set).
+		if flags.Host != "" {
+			cfg.Host = flags.Host
+		}
+		if flags.Profile != "" {
+			cfg.Profile = flags.Profile
+		}
+
+		// Verify that the host is valid (if set).
+		if cfg.Host != "" {
+			err = validateHost(cfg.Host)
+			if err != nil {
+				return err
+			}
+		}
+
+		ctx := cmd.Context()
+		if cmdio.IsInTTY(ctx) && cmdio.IsOutTTY(ctx) {
+			err = configureInteractive(cmd, &flags, &cfg)
+		} else {
+			err = configureNonInteractive(cmd, &flags, &cfg)
 		}
-		err = fn(cmd, ctx, &cfg)
 		if err != nil {
 			return err
 		}
@@ -161,9 +135,10 @@ func newConfigureCommand() *cobra.Command {
 
 		// Save profile to config file.
 		return databrickscfg.SaveToProfile(ctx, &config.Config{
-			Profile: cfg.Profile,
-			Host:    cfg.Host,
-			Token:   cfg.Token,
+			Profile:   cfg.Profile,
+			Host:      cfg.Host,
+			Token:     cfg.Token,
+			ClusterID: cfg.ClusterID,
 		})
 	}
 

cmd/configure/flags.go

@@ -0,0 +1,25 @@
+package configure
+
+import (
+	"github.com/spf13/cobra"
+)
+
+type configureFlags struct {
+	Host    string
+	Profile string
+
+	// Flag to request a prompt for cluster configuration.
+	ConfigureCluster bool
+}
+
+// Register flags with command.
+func (f *configureFlags) Register(cmd *cobra.Command) {
+	cmd.Flags().StringVar(&f.Host, "host", "", "Databricks workspace host.")
+	cmd.Flags().StringVar(&f.Profile, "profile", "DEFAULT", "Name for the connection profile to configure.")
+	cmd.Flags().BoolVar(&f.ConfigureCluster, "configure-cluster", false, "Prompts to configure cluster")
+
+	// Include token flag for compatibility with the legacy CLI.
+	// It doesn't actually do anything because we always use PATs.
+	cmd.Flags().Bool("token", true, "Configure using Databricks Personal Access Token")
+	cmd.Flags().MarkHidden("token")
+}

cmd/configure/host.go

@@ -0,0 +1,20 @@
+package configure
+
+import (
+	"fmt"
+	"net/url"
+)
+
+func validateHost(s string) error {
+	u, err := url.Parse(s)
+	if err != nil {
+		return err
+	}
+	if u.Host == "" || u.Scheme != "https" {
+		return fmt.Errorf("must start with https://")
+	}
+	if u.Path != "" && u.Path != "/" {
+		return fmt.Errorf("must use empty path")
+	}
+	return nil
+}

cmd/configure/host_test.go

@@ -0,0 +1,29 @@
+package configure
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValidateHost(t *testing.T) {
+	var err error
+
+	// Must start with https://
+	err = validateHost("/path")
+	assert.ErrorContains(t, err, "must start with https://")
+	err = validateHost("http://host")
+	assert.ErrorContains(t, err, "must start with https://")
+	err = validateHost("ftp://host")
+
+	// Must use empty path
+	assert.ErrorContains(t, err, "must start with https://")
+	err = validateHost("https://host/path")
+	assert.ErrorContains(t, err, "must use empty path")
+
+	// Ignore query params
+	err = validateHost("https://host/?query")
+	assert.NoError(t, err)
+	err = validateHost("https://host/")
+	assert.NoError(t, err)
+}