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://<HOST>.staging.cloud.databricks.com
{
  "access_token": "...",
  "token_type": "Bearer",
  "expiry": "2023-07-11T12:56:59.929671+02:00"
}
$ ./cli auth login
Databricks Host: https://<HOST>.staging.cloud.databricks.com
Databricks Profile Name: <HOST>-test
Profile <HOST>-test was successfully saved
$ ./cli auth login
Databricks Host: https://accounts.cloud.databricks.com
Databricks Account ID: <ACCOUNTID>
Databricks Profile Name: ACCOUNT-<ACCOUNTID>-test
Profile ACCOUNT-<ACCOUNTID>-test was successfully saved 
```

---------

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>

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

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 {

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
 		}