From d12877a350265ce991065da5457813b0d4741227 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 15:03:38 +0100 Subject: [PATCH 1/5] Add --configure-cluster flag to configure command --- cmd/configure/configure.go | 124 +++++++++++++++--------------------- cmd/configure/flags.go | 38 +++++++++++ cmd/configure/flags_test.go | 1 + cmd/configure/host.go | 20 ++++++ cmd/configure/host_test.go | 29 +++++++++ 5 files changed, 139 insertions(+), 73 deletions(-) create mode 100644 cmd/configure/flags.go create mode 100644 cmd/configure/flags_test.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 55ede5385b..48054e3b6f 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -1,65 +1,18 @@ 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 == "" { @@ -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 precence 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,31 @@ func newConfigureCommand() *cobra.Command { return fmt.Errorf("unable to instantiate configuration from environment variables: %w", err) } + // 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() interactive := cmdio.IsInTTY(ctx) && cmdio.IsOutTTY(ctx) - var fn func(*cobra.Command, context.Context, *config.Config) error + var fn func(*cobra.Command, *configureFlags, *config.Config) error if interactive { fn = configureInteractive } else { fn = configureNonInteractive } - err = fn(cmd, ctx, &cfg) + err = fn(cmd, &flags, &cfg) if err != nil { return err } @@ -161,9 +138,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 0000000000..2c66a76691 --- /dev/null +++ b/cmd/configure/flags.go @@ -0,0 +1,38 @@ +package configure + +import ( + "github.com/databricks/databricks-sdk-go/config" + "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") +} + +func (f *configureFlags) PopulateConfig(cfg *config.Config) error { + if f.Host != "" { + cfg.Host = f.Host + } + + if f.Profile != "" { + cfg.Profile = f.Profile + } + + return nil +} diff --git a/cmd/configure/flags_test.go b/cmd/configure/flags_test.go new file mode 100644 index 0000000000..7fa1b68a95 --- /dev/null +++ b/cmd/configure/flags_test.go @@ -0,0 +1 @@ +package configure diff --git a/cmd/configure/host.go b/cmd/configure/host.go new file mode 100644 index 0000000000..781c123871 --- /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 0000000000..a4af199d69 --- /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) +} From 236a8db4034eb326effca1c6792cb2c83d204911 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 15:05:05 +0100 Subject: [PATCH 2/5] Lowercase --- cmd/configure/configure.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index 48054e3b6f..2533d4efda 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -17,7 +17,7 @@ func configureInteractive(cmd *cobra.Command, flags *configureFlags, cfg *config // 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 @@ -31,7 +31,7 @@ func configureInteractive(cmd *cobra.Command, flags *configureFlags, cfg *config // 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 { From bfb55c41194c84b982ee586aaea903f6c8e86b99 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 15:06:09 +0100 Subject: [PATCH 3/5] Cleanup --- cmd/configure/flags.go | 13 ------------- cmd/configure/flags_test.go | 1 - 2 files changed, 14 deletions(-) delete mode 100644 cmd/configure/flags_test.go diff --git a/cmd/configure/flags.go b/cmd/configure/flags.go index 2c66a76691..80e6502689 100644 --- a/cmd/configure/flags.go +++ b/cmd/configure/flags.go @@ -1,7 +1,6 @@ package configure import ( - "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" ) @@ -24,15 +23,3 @@ func (f *configureFlags) Register(cmd *cobra.Command) { cmd.Flags().Bool("token", true, "Configure using Databricks Personal Access Token") cmd.Flags().MarkHidden("token") } - -func (f *configureFlags) PopulateConfig(cfg *config.Config) error { - if f.Host != "" { - cfg.Host = f.Host - } - - if f.Profile != "" { - cfg.Profile = f.Profile - } - - return nil -} diff --git a/cmd/configure/flags_test.go b/cmd/configure/flags_test.go deleted file mode 100644 index 7fa1b68a95..0000000000 --- a/cmd/configure/flags_test.go +++ /dev/null @@ -1 +0,0 @@ -package configure From cf3fd6c36c2b24bd1c3e43ca04f3ee8a0cab00f8 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 20:48:39 +0100 Subject: [PATCH 4/5] Update cmd/configure/configure.go Co-authored-by: Miles Yucht --- cmd/configure/configure.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index 2533d4efda..8011d0a6fa 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -61,7 +61,7 @@ func configureNonInteractive(cmd *cobra.Command, flags *configureFlags, cfg *con return fmt.Errorf("host must be set in non-interactive mode") } - // Check precence of cluster ID before reading token to fail fast. + // 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") } From 8af283456c8b1cb34043642b3338a25305c97594 Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Thu, 23 Nov 2023 20:51:04 +0100 Subject: [PATCH 5/5] Inline calls --- cmd/configure/configure.go | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index 8011d0a6fa..1c4d2e6b27 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -120,14 +120,11 @@ The host must be specified with the --host flag or the DATABRICKS_HOST environme } ctx := cmd.Context() - interactive := cmdio.IsInTTY(ctx) && cmdio.IsOutTTY(ctx) - var fn func(*cobra.Command, *configureFlags, *config.Config) error - if interactive { - fn = configureInteractive + if cmdio.IsInTTY(ctx) && cmdio.IsOutTTY(ctx) { + err = configureInteractive(cmd, &flags, &cfg) } else { - fn = configureNonInteractive + err = configureNonInteractive(cmd, &flags, &cfg) } - err = fn(cmd, &flags, &cfg) if err != nil { return err }