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

Update configure command (#482) 

## Changes

This now uses:
* libs/cmdio to determine interactivity and perform prompting
* libs/databrickscfg to persist the profile

It loads a config.Config structure from the environment just like we do
for unified authentication. It is therefore possible to specify both the
host and token with environment variables.

## Tests

```
pieter.noordhuis@L4GHXDT29P /tmp % export DATABRICKS_CONFIG_FILE=.databrickscfg
pieter.noordhuis@L4GHXDT29P /tmp % databricks configure
Databricks Host: https://foo.bar
Personal Access Token: *****
pieter.noordhuis@L4GHXDT29P /tmp % cat .databrickscfg
[DEFAULT]
host  = https://foo.bar
token = token
pieter.noordhuis@L4GHXDT29P /tmp % echo token | databricks configure
Error: host must be set in non-interactive mode
pieter.noordhuis@L4GHXDT29P /tmp % echo token | databricks configure --host foo
Error: must start with https://
pieter.noordhuis@L4GHXDT29P /tmp % echo token | databricks configure --host https://foo
pieter.noordhuis@L4GHXDT29P /tmp % cat .databrickscfg
[DEFAULT]
host  = https://foo
token = token
pieter.noordhuis@L4GHXDT29P /tmp % cat .databrickscfg
pieter.noordhuis@L4GHXDT29P /tmp % databricks configure --host https://foo
Personal Access Token: ******
pieter.noordhuis@L4GHXDT29P /tmp % cat .databrickscfg
[DEFAULT]
host  = https://foo
token = token2
```
This commit is contained in:
Pieter Noordhuis 2023-06-15 14:50:19 +02:00 committed by GitHub
parent becf4d9c31
commit b9406efd27
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 178 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

225
cmd/configure/configure.go
View File

@ -1,124 +1,167 @@
package configure
import (
"bytes"
"errors"
"context"
"fmt"
"os"
"path/filepath"
"net/url"
"github.com/databricks/cli/cmd/root"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/cli/libs/databrickscfg"
"github.com/databricks/databricks-sdk-go/config"
"github.com/spf13/cobra"
"gopkg.in/ini.v1"
)
type Configs struct {
Host string `ini:"host"`
Token string `ini:"token,omitempty"`
Profile string `ini:"-"`
}
var tokenMode bool
func (cfg *Configs) loadNonInteractive(cmd *cobra.Command) error {
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)
func validateHost(s string) error {
u, err := url.Parse(s)
if err != nil {
return err
}
if n != 1 {
return fmt.Errorf("exactly 1 argument required")
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
}
// Ask user to specify the host if not already set.
if cfg.Host == "" {
prompt := cmdio.Prompt(ctx)
prompt.Label = "Databricks Host"
prompt.Default = "https://"
prompt.AllowEdit = true
prompt.Validate = validateHost
out, err := prompt.Run()
if err != nil {
return err
}
cfg.Host = out
}
// Ask user to specify the token is not already set.
if cfg.Token == "" {
prompt := cmdio.Prompt(ctx)
prompt.Label = "Personal Access Token"
prompt.Mask = '*'
out, err := prompt.Run()
if err != nil {
return err
}
cfg.Token = out
}
return nil
}
func configureNonInteractive(cmd *cobra.Command, ctx context.Context, cfg *config.Config) error {
err := configureFromFlags(cmd, ctx, cfg)
if err != nil {
return err
}
if cfg.Host == "" {
return fmt.Errorf("host 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)
if err != nil {
return err
}
}
return nil
}
var configureCmd = &cobra.Command{
Use: "configure",
Short: "Configure authentication",
Use: "configure",
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.
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.
`,
Hidden: true,
RunE: func(cmd *cobra.Command, args []string) error {
profile, err := cmd.Flags().GetString("profile")
var cfg config.Config
// Load environment variables, possibly the DEFAULT profile.
err := config.ConfigAttributes.Configure(&cfg)
if err != nil {
return fmt.Errorf("read --profile flag: %w", err)
return fmt.Errorf("unable to instantiate configuration from environment variables: %w", err)
}
path := os.Getenv("DATABRICKS_CONFIG_FILE")
if path == "" {
path, err = os.UserHomeDir()
if err != nil {
return fmt.Errorf("homedir: %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
}
if filepath.Base(path) == ".databrickscfg" {
path = filepath.Dir(path)
}
err = os.MkdirAll(path, os.ModeDir|os.ModePerm)
err = fn(cmd, ctx, &cfg)
if err != nil {
return fmt.Errorf("create config dir: %w", err)
}
cfgPath := filepath.Join(path, ".databrickscfg")
_, err = os.Stat(cfgPath)
if errors.Is(err, os.ErrNotExist) {
file, err := os.Create(cfgPath)
if err != nil {
return fmt.Errorf("create config file: %w", err)
}
file.Close()
} else if err != nil {
return fmt.Errorf("open config file: %w", err)
return err
}
ini_cfg, err := ini.Load(cfgPath)
if err != nil {
return fmt.Errorf("load config file: %w", err)
}
cfg := &Configs{"", "", profile}
err = ini_cfg.Section(profile).MapTo(cfg)
if err != nil {
return fmt.Errorf("unmarshal loaded config: %w", err)
}
// Clear the Databricks CLI path in token mode.
// This is relevant for OAuth only.
cfg.DatabricksCliPath = ""
err = cfg.loadNonInteractive(cmd)
if err != nil {
return fmt.Errorf("reading configs: %w", err)
}
err = ini_cfg.Section(profile).ReflectFrom(cfg)
if err != nil {
return fmt.Errorf("marshall config: %w", err)
}
var buffer bytes.Buffer
if ini_cfg.Section("DEFAULT").Body() != "" {
//This configuration makes the ini library write the DEFAULT header explicitly.
//DEFAULT section might be empty
ini.DefaultHeader = true
}
_, err = ini_cfg.WriteTo(&buffer)
if err != nil {
return fmt.Errorf("write config to buffer: %w", err)
}
err = os.WriteFile(cfgPath, buffer.Bytes(), os.ModePerm)
if err != nil {
return fmt.Errorf("write congfig to file: %w", err)
}
return nil
// Save profile to config file.
return databrickscfg.SaveToProfile(ctx, &cfg)
},
}
func init() {
root.RootCmd.AddCommand(configureCmd)
configureCmd.Flags().BoolVarP(&tokenMode, "token", "t", false, "Configure using Databricks Personal Access Token")
configureCmd.Flags().String("host", "", "Host to connect to.")
configureCmd.Flags().String("profile", "DEFAULT", "CLI connection profile to use.")
configureCmd.Flags().String("host", "", "Databricks workspace host.")
configureCmd.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.
configureCmd.Flags().BoolP("token", "t", true, "Configure using Databricks Personal Access Token")
configureCmd.Flags().MarkHidden("token")
}

19
cmd/configure/configure_test.go
View File

@ -52,7 +52,7 @@ func TestDefaultConfigureNoInteractive(t *testing.T) {
})
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "host"})
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "https://host"})
err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err)
@ -67,7 +67,7 @@ func TestDefaultConfigureNoInteractive(t *testing.T) {
defaultSection, err := cfg.GetSection("DEFAULT")
assert.NoError(t, err)
assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "host", "https://host")
assertKeyValueInSection(t, defaultSection, "token", "token")
}
@ -75,8 +75,8 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
//TODO: Replace with similar test code from go SDK, once we start using it directly
ctx := context.Background()
tempHomeDir := setup(t)
cfgFileDir := filepath.Join(tempHomeDir, "test")
t.Setenv("DATABRICKS_CONFIG_FILE", cfgFileDir)
cfgPath := filepath.Join(tempHomeDir, ".databrickscfg")
t.Setenv("DATABRICKS_CONFIG_FILE", cfgPath)
inp := getTempFileWithContent(t, tempHomeDir, "token\n")
defer inp.Close()
@ -84,12 +84,11 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "host"})
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "https://host"})
err := root.RootCmd.ExecuteContext(ctx)
assert.NoError(t, err)
cfgPath := filepath.Join(cfgFileDir, ".databrickscfg")
_, err = os.Stat(cfgPath)
assert.NoError(t, err)
@ -99,25 +98,25 @@ func TestConfigFileFromEnvNoInteractive(t *testing.T) {
defaultSection, err := cfg.GetSection("DEFAULT")
assert.NoError(t, err)
assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "host", "https://host")
assertKeyValueInSection(t, defaultSection, "token", "token")
}
func TestCustomProfileConfigureNoInteractive(t *testing.T) {
ctx := context.Background()
tempHomeDir := setup(t)
cfgPath := filepath.Join(tempHomeDir, ".databrickscfg")
inp := getTempFileWithContent(t, tempHomeDir, "token\n")
defer inp.Close()
oldStdin := os.Stdin
t.Cleanup(func() { os.Stdin = oldStdin })
os.Stdin = inp
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "host", "--profile", "CUSTOM"})
root.RootCmd.SetArgs([]string{"configure", "--token", "--host", "https://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)
@ -127,6 +126,6 @@ func TestCustomProfileConfigureNoInteractive(t *testing.T) {
defaultSection, err := cfg.GetSection("CUSTOM")
assert.NoError(t, err)
assertKeyValueInSection(t, defaultSection, "host", "host")
assertKeyValueInSection(t, defaultSection, "host", "https://host")
assertKeyValueInSection(t, defaultSection, "token", "token")
}

36
libs/cmdio/io.go
View File

@ -51,7 +51,7 @@ func IsInteractive(ctx context.Context) bool {
}
// IsTTY detects if io.Writer is a terminal.
func IsTTY(w io.Writer) bool {
func IsTTY(w any) bool {
f, ok := w.(*os.File)
if !ok {
return false
@ -60,6 +60,24 @@ func IsTTY(w io.Writer) bool {
return isatty.IsTerminal(fd) || isatty.IsCygwinTerminal(fd)
}
// IsInTTY detects if the input reader is a terminal.
func IsInTTY(ctx context.Context) bool {
c := fromContext(ctx)
return IsTTY(c.in)
}
// IsOutTTY detects if the output writer is a terminal.
func IsOutTTY(ctx context.Context) bool {
c := fromContext(ctx)
return IsTTY(c.out)
}
// IsErrTTY detects if the error writer is a terminal.
func IsErrTTY(ctx context.Context) bool {
c := fromContext(ctx)
return IsTTY(c.err)
}
// IsTTY detects if stdout is a terminal. It assumes that stderr is terminal as well
func (c *cmdIO) IsTTY() bool {
f, ok := c.out.(*os.File)
@ -170,6 +188,22 @@ func Secret(ctx context.Context) (value string, err error) {
return c.Secret()
}
type nopWriteCloser struct {
io.Writer
}
func (nopWriteCloser) Close() error {
return nil
}
func Prompt(ctx context.Context) *promptui.Prompt {
c := fromContext(ctx)
return &promptui.Prompt{
Stdin: io.NopCloser(c.in),
Stdout: nopWriteCloser{c.out},
}
}
func (c *cmdIO) Spinner(ctx context.Context) chan string {
var sp *spinner.Spinner
if c.interactive {