From 9a0888126c8a30433a1d40d710e1162bc02a7bd2 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Wed, 12 Jul 2023 17:36:09 +0200 Subject: [PATCH] Improve auth login experience (#570) ## Changes Currently, `databricks auth login` is difficult to use. If a user types this command in, the command fails with ``` Error: init: cannot fetch credentials ``` after prompting for a profile name. To make this experience smoother, this change ensures that the host, and if necessary, the account ID, are prompted for input from the user if they aren't provided on the CLI. ## Tests Manual tests: ``` $ ./cli auth token Databricks Host: https://.staging.cloud.databricks.com { "access_token": "...", "token_type": "Bearer", "expiry": "2023-07-11T12:56:59.929671+02:00" } $ ./cli auth login Databricks Host: https://.staging.cloud.databricks.com Databricks Profile Name: -test Profile -test was successfully saved $ ./cli auth login Databricks Host: https://accounts.cloud.databricks.com Databricks Account ID: Databricks Profile Name: ACCOUNT--test Profile ACCOUNT--test was successfully saved ``` --------- Co-authored-by: Pieter Noordhuis --- cmd/auth/auth.go | 35 ++++++++++++++++++++++++++++++--- cmd/auth/login.go | 50 ++++++++++++++++++++++++++++++++++------------- 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 3efaca57..b7e8d2d7 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,36 @@ var authCmd = &cobra.Command{ Short: "Authentication related commands", } -var perisistentAuth auth.PersistentAuth +var persistentAuth auth.PersistentAuth + +func promptForHost(ctx context.Context) (string, 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 + } + return host, nil +} + +func promptForAccountID(ctx context.Context) (string, 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 + } + return accountId, 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 6b708e95..37d44c08 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -17,16 +17,46 @@ 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 + } + + host, err := promptForHost(ctx) + if err != nil { + return err + } + persistentAuth.Host = host + 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() - defer perisistentAuth.Close() - ctx, cancel := context.WithTimeout(cmd.Context(), loginTimeout) + // 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 == "" { + accountId, err := promptForAccountID(ctx) + if err != nil { + return err + } + persistentAuth.AccountID = accountId + } + cfg.AccountID = persistentAuth.AccountID + + ctx, cancel := context.WithTimeout(ctx, loginTimeout) defer cancel() var profileName string @@ -36,7 +66,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 +74,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 f2754fa6..1b8d8b13 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 }