Improve auth login experience

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

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,27 +64,19 @@ 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 {
 				return err
 			}
 			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
 		}