From d985601d3051657c40541e0448c1574e2cb353f7 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 20:56:48 +0100 Subject: [PATCH] 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 --- cmd/configure/configure.go | 137 +++++++++++++++---------------------- cmd/configure/flags.go | 25 +++++++ cmd/configure/host.go | 20 ++++++ cmd/configure/host_test.go | 29 ++++++++ 4 files changed, 130 insertions(+), 81 deletions(-) create mode 100644 cmd/configure/flags.go create mode 100644 cmd/configure/host.go create mode 100644 cmd/configure/host_test.go diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index 55ede538..1c4d2e6b 100644 --- a/cmd/configure/configure.go +++ b/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, }) } diff --git a/cmd/configure/flags.go b/cmd/configure/flags.go new file mode 100644 index 00000000..80e65026 --- /dev/null +++ b/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") +} diff --git a/cmd/configure/host.go b/cmd/configure/host.go new file mode 100644 index 00000000..781c1238 --- /dev/null +++ b/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 +} diff --git a/cmd/configure/host_test.go b/cmd/configure/host_test.go new file mode 100644 index 00000000..a4af199d --- /dev/null +++ b/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) +}