steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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>
This commit is contained in:
Pieter Noordhuis 2023-11-23 20:56:48 +01:00 committed by GitHub
parent 07c4c90772
commit d985601d30
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 130 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

137
cmd/configure/configure.go
View File

@ -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,
})
}

25
cmd/configure/flags.go Normal file
View File

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

20
cmd/configure/host.go Normal file
View File

@ -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
}

29
cmd/configure/host_test.go Normal file
View File

@ -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)
}