From 04fbf25aa20f1a19405b860846e9e1c73995bbc1 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Tue, 11 Jul 2023 11:58:12 +0200 Subject: [PATCH 1/4] Improve auth login experience --- cmd/auth/auth.go | 37 +++++++++++++++++++++++++++++++++--- cmd/auth/login.go | 48 +++++++++++++++++++++++++++++++++-------------- cmd/auth/token.go | 12 +++++++----- 3 files changed, 75 insertions(+), 22 deletions(-) diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 3efaca5727..0bdb5360be 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -1,8 +1,11 @@ package auth import ( + "context" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdio" "github.com/spf13/cobra" ) @@ -11,10 +14,38 @@ var authCmd = &cobra.Command{ Short: "Authentication related commands", } -var perisistentAuth auth.PersistentAuth +var persistentAuth auth.PersistentAuth + +func promptForHost(ctx context.Context) error { + prompt := cmdio.Prompt(ctx) + prompt.Label = "Databricks Host" + prompt.Default = "https://" + prompt.AllowEdit = true + // Validate? + host, err := prompt.Run() + if err != nil { + return err + } + persistentAuth.Host = host + return nil +} + +func promptForAccountId(ctx context.Context) error { + prompt := cmdio.Prompt(ctx) + prompt.Label = "Databricks Account ID" + prompt.Default = "" + prompt.AllowEdit = true + // Validate? + accountId, err := prompt.Run() + if err != nil { + return err + } + persistentAuth.AccountID = accountId + return nil +} func init() { root.RootCmd.AddCommand(authCmd) - authCmd.PersistentFlags().StringVar(&perisistentAuth.Host, "host", perisistentAuth.Host, "Databricks Host") - authCmd.PersistentFlags().StringVar(&perisistentAuth.AccountID, "account-id", perisistentAuth.AccountID, "Databricks Account ID") + authCmd.PersistentFlags().StringVar(&persistentAuth.Host, "host", persistentAuth.Host, "Databricks Host") + authCmd.PersistentFlags().StringVar(&persistentAuth.AccountID, "account-id", persistentAuth.AccountID, "Databricks Account ID") } diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 6b708e9575..dd51a5b90a 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -17,16 +17,44 @@ import ( var loginTimeout time.Duration var configureCluster bool +func configureHost(ctx context.Context, args []string, argIndex int) error { + if len(args) > argIndex { + persistentAuth.Host = args[argIndex] + return nil + } + + err := promptForHost(ctx) + if err != nil { + return err + } + return nil +} + var loginCmd = &cobra.Command{ Use: "login [HOST]", Short: "Authenticate this machine", RunE: func(cmd *cobra.Command, args []string) error { - if perisistentAuth.Host == "" && len(args) == 1 { - perisistentAuth.Host = args[0] + ctx := cmd.Context() + if persistentAuth.Host == "" { + configureHost(ctx, args, 0) + } + defer persistentAuth.Close() + + // We need the config without the profile before it's used to initialise new workspace client below. + // Otherwise it will complain about non existing profile because it was not yet saved. + cfg := config.Config{ + Host: persistentAuth.Host, + AuthType: "databricks-cli", } + if cfg.IsAccountClient() && persistentAuth.AccountID == "" { + err := promptForAccountId(ctx) + if err != nil { + return err + } + } + cfg.AccountID = persistentAuth.AccountID - defer perisistentAuth.Close() - ctx, cancel := context.WithTimeout(cmd.Context(), loginTimeout) + ctx, cancel := context.WithTimeout(ctx, loginTimeout) defer cancel() var profileName string @@ -36,7 +64,7 @@ var loginCmd = &cobra.Command{ } else { prompt := cmdio.Prompt(ctx) prompt.Label = "Databricks Profile Name" - prompt.Default = perisistentAuth.ProfileName() + prompt.Default = persistentAuth.ProfileName() prompt.AllowEdit = true profile, err := prompt.Run() if err != nil { @@ -44,19 +72,11 @@ var loginCmd = &cobra.Command{ } profileName = profile } - err := perisistentAuth.Challenge(ctx) + err := persistentAuth.Challenge(ctx) if err != nil { return err } - // We need the config without the profile before it's used to initialise new workspace client below. - // Otherwise it will complain about non existing profile because it was not yet saved. - cfg := config.Config{ - Host: perisistentAuth.Host, - AccountID: perisistentAuth.AccountID, - AuthType: "databricks-cli", - } - if configureCluster { w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg)) if err != nil { diff --git a/cmd/auth/token.go b/cmd/auth/token.go index f2754fa69d..1b8d8b131b 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -15,13 +15,15 @@ var tokenCmd = &cobra.Command{ Use: "token [HOST]", Short: "Get authentication token", RunE: func(cmd *cobra.Command, args []string) error { - if perisistentAuth.Host == "" && len(args) == 1 { - perisistentAuth.Host = args[0] + ctx := cmd.Context() + if persistentAuth.Host == "" { + configureHost(ctx, args, 0) } - defer perisistentAuth.Close() - ctx, cancel := context.WithTimeout(cmd.Context(), tokenTimeout) + defer persistentAuth.Close() + + ctx, cancel := context.WithTimeout(ctx, tokenTimeout) defer cancel() - t, err := perisistentAuth.Load(ctx) + t, err := persistentAuth.Load(ctx) if err != nil { return err } From ac8b3b7b4fde798cf048a60e204403d4805b7998 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Tue, 11 Jul 2023 18:09:37 +0200 Subject: [PATCH 2/4] address comments --- cmd/auth/auth.go | 14 ++++++-------- cmd/auth/login.go | 6 ++++-- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 0bdb5360be..5cd9a5f211 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -16,7 +16,7 @@ var authCmd = &cobra.Command{ var persistentAuth auth.PersistentAuth -func promptForHost(ctx context.Context) error { +func promptForHost(ctx context.Context) (string, error) { prompt := cmdio.Prompt(ctx) prompt.Label = "Databricks Host" prompt.Default = "https://" @@ -24,13 +24,12 @@ func promptForHost(ctx context.Context) error { // Validate? host, err := prompt.Run() if err != nil { - return err + return "", err } - persistentAuth.Host = host - return nil + return host, nil } -func promptForAccountId(ctx context.Context) error { +func promptForAccountId(ctx context.Context) (string, error) { prompt := cmdio.Prompt(ctx) prompt.Label = "Databricks Account ID" prompt.Default = "" @@ -38,10 +37,9 @@ func promptForAccountId(ctx context.Context) error { // Validate? accountId, err := prompt.Run() if err != nil { - return err + return "", err } - persistentAuth.AccountID = accountId - return nil + return accountId, nil } func init() { diff --git a/cmd/auth/login.go b/cmd/auth/login.go index dd51a5b90a..deae4e7b0a 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -23,10 +23,11 @@ func configureHost(ctx context.Context, args []string, argIndex int) error { return nil } - err := promptForHost(ctx) + host, err := promptForHost(ctx) if err != nil { return err } + persistentAuth.Host = host return nil } @@ -47,10 +48,11 @@ var loginCmd = &cobra.Command{ AuthType: "databricks-cli", } if cfg.IsAccountClient() && persistentAuth.AccountID == "" { - err := promptForAccountId(ctx) + accountId, err := promptForAccountId(ctx) if err != nil { return err } + persistentAuth.AccountID = accountId } cfg.AccountID = persistentAuth.AccountID From 2b8872d5e17522beab78faf8f83480f557b11651 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Wed, 12 Jul 2023 17:27:05 +0200 Subject: [PATCH 3/4] Update cmd/auth/auth.go Co-authored-by: Pieter Noordhuis --- cmd/auth/auth.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 5cd9a5f211..b7e8d2d788 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -29,7 +29,7 @@ func promptForHost(ctx context.Context) (string, error) { return host, nil } -func promptForAccountId(ctx context.Context) (string, error) { +func promptForAccountID(ctx context.Context) (string, error) { prompt := cmdio.Prompt(ctx) prompt.Label = "Databricks Account ID" prompt.Default = "" From 4793be3f14f76188132914deaef93a0dbd04d4a3 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Wed, 12 Jul 2023 17:30:29 +0200 Subject: [PATCH 4/4] fix --- cmd/auth/login.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index deae4e7b0a..37d44c0848 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -48,7 +48,7 @@ var loginCmd = &cobra.Command{ AuthType: "databricks-cli", } if cfg.IsAccountClient() && persistentAuth.AccountID == "" { - accountId, err := promptForAccountId(ctx) + accountId, err := promptForAccountID(ctx) if err != nil { return err }