From f41eff40d99595f278730c37109b4d8b64472020 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 16:50:46 +0200
Subject: [PATCH 01/20] add flag for --profile

---
 cmd/auth/login.go | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index b0bc7a853b..b18cd85e33 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -39,19 +39,17 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 
 	var loginTimeout time.Duration
 	var configureCluster bool
+	var profileName string
 	cmd.Flags().DurationVar(&loginTimeout, "timeout", auth.DefaultTimeout,
 		"Timeout for completing login challenge in the browser")
 	cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
 		"Prompts to configure cluster")
+	cmd.Flags().StringVarP(&profileName, "profile", "p", "", `Name of the profile.`)
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 
-		var profileName string
-		profileFlag := cmd.Flag("profile")
-		if profileFlag != nil && profileFlag.Value.String() != "" {
-			profileName = profileFlag.Value.String()
-		} else if cmdio.IsInTTY(ctx) {
+		if profileName == "" && cmdio.IsInTTY(ctx) {
 			prompt := cmdio.Prompt(ctx)
 			prompt.Label = "Databricks Profile Name"
 			prompt.Default = persistentAuth.ProfileName()

From c7a7696b2fe9dfcb4e9ad84a92f3e9034612e6d3 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:37:49 +0200
Subject: [PATCH 02/20] Improvements to auth login command

---
 cmd/auth/auth.go   |  27 ----------
 cmd/auth/login.go  | 131 +++++++++++++++++++++++++++++++++------------
 libs/auth/oauth.go |   3 +-
 3 files changed, 99 insertions(+), 62 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 4af2a7a71c..c8e1118705 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -1,10 +1,7 @@
 package auth
 
 import (
-	"context"
-
 	"github.com/databricks/cli/libs/auth"
-	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
@@ -32,27 +29,3 @@ GCP: https://docs.gcp.databricks.com/en/dev-tools/auth/index.html`,
 	cmd.AddCommand(newDescribeCommand())
 	return cmd
 }
-
-func promptForHost(ctx context.Context) (string, error) {
-	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Host (e.g. https://<databricks-instance>.cloud.databricks.com)"
-	// 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
-}
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index b18cd85e33..0510e81a8f 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
@@ -15,17 +16,101 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func configureHost(ctx context.Context, persistentAuth *auth.PersistentAuth, args []string, argIndex int) error {
-	if len(args) > argIndex {
-		persistentAuth.Host = args[argIndex]
+func promptForHost(ctx context.Context) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify a host using --host")
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks Host"
+	prompt.Validate = func(host string) error {
+		if !strings.HasPrefix(host, "https://") {
+			return fmt.Errorf("host URL must have a https:// prefix")
+		}
 		return nil
 	}
+	return prompt.Run()
+}
+
+func promptForProfile(ctx context.Context, dv string) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		return "", errors.New("the command is being run in a non-interactive environment, please specify a profile using --profile")
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks Profile Name"
+	prompt.Default = dv
+	prompt.AllowEdit = true
+	return prompt.Run()
+}
+
+func promptForAccountID(ctx context.Context) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		return "", errors.New("the command is being run in a non-interactive environment, please specify an account ID using --account-id")
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks Account ID"
+	prompt.Default = ""
+	prompt.AllowEdit = true
+	return prompt.Run()
+}
 
-	host, err := promptForHost(ctx)
+func getHostFromProfile(ctx context.Context, profileName string) (string, error) {
+	_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
+		return p.Name == profileName
+	})
+
+	// Tolerate ErrNoConfiguration here, as we will write out a configuration file
+	// as part of the login flow.
+	// TODO: verify that a configuration file is written out as part of the login flow.
+	if err != nil && errors.Is(err, databrickscfg.ErrNoConfiguration) {
+		return "", nil
+	}
 	if err != nil {
-		return err
+		return "", err
+	}
+
+	// Return host from profile
+	if len(profiles) > 0 && profiles[0].Host != "" {
+		return profiles[0].Host, nil
+	}
+	return "", nil
+}
+
+// Sets the host in the persistentAuth object based on the provided arguments and flags.
+// Follows the following precedence:
+// 1. [HOST] (first positional argument) or --host flag. Error if both are specified.
+// 2. Profile host, if available.
+// 3. Prompt the user for the host.
+func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
+	// If both [HOST] and --host are provided, return an error.
+	// TODO: test for error for this case.
+	if len(args) > 0 && persistentAuth.Host != "" {
+		return fmt.Errorf("please only provide a host as an argument or a flag, not both")
+	}
+
+	// If [HOST] is provided, set the host to the provided positional argument.
+	if len(args) > 0 && persistentAuth.Host == "" {
+		persistentAuth.Host = args[0]
+	}
+
+	// If neither [HOST] nor --host are provided, and the profile has a host, use it.
+	// Otherwise, prompt the user for a host.
+	if len(args) == 0 && persistentAuth.Host == "" {
+		hostName, err := getHostFromProfile(ctx, profileName)
+		if err != nil {
+			return err
+		}
+		if hostName == "" {
+			var err error
+			hostName, err = promptForHost(ctx)
+			if err != nil {
+				return err
+			}
+		}
+		persistentAuth.Host = hostName
 	}
-	persistentAuth.Host = host
 	return nil
 }
 
@@ -49,23 +134,21 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 
-		if profileName == "" && cmdio.IsInTTY(ctx) {
-			prompt := cmdio.Prompt(ctx)
-			prompt.Label = "Databricks Profile Name"
-			prompt.Default = persistentAuth.ProfileName()
-			prompt.AllowEdit = true
-			profile, err := prompt.Run()
+		// If the user has not specified a profile name, and we are in an interactive
+		// session, prompt for one.
+		if profileName == "" {
+			var err error
+			profileName, err = promptForProfile(ctx, persistentAuth.DefaultProfileName())
 			if err != nil {
 				return err
 			}
-			profileName = profile
 		}
 
+		// Set the host based on the provided arguments and flags.
 		err := setHost(ctx, profileName, persistentAuth, args)
 		if err != nil {
 			return err
 		}
-		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.
@@ -89,6 +172,7 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 		if err != nil {
 			return err
 		}
+		defer persistentAuth.Close()
 
 		if configureCluster {
 			w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
@@ -124,22 +208,3 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 
 	return cmd
 }
-
-func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
-	// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
-	_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
-		return p.Name == profileName
-	})
-	// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
-	if err != nil && !errors.Is(err, databrickscfg.ErrNoConfiguration) {
-		return err
-	}
-	if persistentAuth.Host == "" {
-		if len(profiles) > 0 && profiles[0].Host != "" {
-			persistentAuth.Host = profiles[0].Host
-		} else {
-			configureHost(ctx, persistentAuth, args, 0)
-		}
-	}
-	return nil
-}
diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go
index 4ce0d4defc..d75d0b683a 100644
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@@ -91,8 +91,7 @@ func (a *PersistentAuth) Load(ctx context.Context) (*oauth2.Token, error) {
 	return refreshed, nil
 }
 
-func (a *PersistentAuth) ProfileName() string {
-	// TODO: get profile name from interactive input
+func (a *PersistentAuth) DefaultProfileName() string {
 	if a.AccountID != "" {
 		return fmt.Sprintf("ACCOUNT-%s", a.AccountID)
 	}

From c1523cd5f52fe84e7ab089e7a6aa5cddb776cbce Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:48:22 +0200
Subject: [PATCH 03/20] test for setHost

---
 cmd/auth/login_test.go           | 35 ++++++++++++++++++++++++++++++++
 cmd/auth/testdata/.databrickscfg |  5 +++++
 2 files changed, 40 insertions(+)
 create mode 100644 cmd/auth/testdata/.databrickscfg

diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go
index 9b834bd0a3..75c58656da 100644
--- a/cmd/auth/login_test.go
+++ b/cmd/auth/login_test.go
@@ -15,3 +15,38 @@ func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
 	err := setHost(ctx, "foo", &auth.PersistentAuth{Host: "test"}, []string{})
 	assert.NoError(t, err)
 }
+
+func TestSetHost(t *testing.T) {
+	ctx := context.Background()
+	var persistentAuth auth.PersistentAuth
+	t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
+
+	// Test error when both flag and argument are provided
+	persistentAuth.Host = "val from --host"
+	err := setHost(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
+
+	// Test setting host from flag
+	persistentAuth.Host = "val from --host"
+	err = setHost(ctx, "profile-1", &persistentAuth, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "val from --host", persistentAuth.Host)
+
+	// Test setting host from argument
+	persistentAuth.Host = ""
+	err = setHost(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	assert.NoError(t, err)
+	assert.Equal(t, "val from [HOST]", persistentAuth.Host)
+
+	// Test setting host from profile
+	persistentAuth.Host = ""
+	err = setHost(ctx, "profile-1", &persistentAuth, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://www.host1.com", persistentAuth.Host)
+
+	// Test setting host from profile
+	persistentAuth.Host = ""
+	err = setHost(ctx, "profile-2", &persistentAuth, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://www.host2.com", persistentAuth.Host)
+}
diff --git a/cmd/auth/testdata/.databrickscfg b/cmd/auth/testdata/.databrickscfg
new file mode 100644
index 0000000000..e815549bc3
--- /dev/null
+++ b/cmd/auth/testdata/.databrickscfg
@@ -0,0 +1,5 @@
+[profile-1]
+host = https://www.host1.com
+
+[profile-2]
+host = https://www.host2.com

From b05eb8b458a76ccf627c60da3ca90c27000a6c77 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:53:59 +0200
Subject: [PATCH 04/20] smaller diff

---
 cmd/auth/login.go | 72 +++++++++++++++++++++++------------------------
 1 file changed, 36 insertions(+), 36 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index 0510e81a8f..55158ab3ae 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -78,42 +78,6 @@ func getHostFromProfile(ctx context.Context, profileName string) (string, error)
 	return "", nil
 }
 
-// Sets the host in the persistentAuth object based on the provided arguments and flags.
-// Follows the following precedence:
-// 1. [HOST] (first positional argument) or --host flag. Error if both are specified.
-// 2. Profile host, if available.
-// 3. Prompt the user for the host.
-func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
-	// If both [HOST] and --host are provided, return an error.
-	// TODO: test for error for this case.
-	if len(args) > 0 && persistentAuth.Host != "" {
-		return fmt.Errorf("please only provide a host as an argument or a flag, not both")
-	}
-
-	// If [HOST] is provided, set the host to the provided positional argument.
-	if len(args) > 0 && persistentAuth.Host == "" {
-		persistentAuth.Host = args[0]
-	}
-
-	// If neither [HOST] nor --host are provided, and the profile has a host, use it.
-	// Otherwise, prompt the user for a host.
-	if len(args) == 0 && persistentAuth.Host == "" {
-		hostName, err := getHostFromProfile(ctx, profileName)
-		if err != nil {
-			return err
-		}
-		if hostName == "" {
-			var err error
-			hostName, err = promptForHost(ctx)
-			if err != nil {
-				return err
-			}
-		}
-		persistentAuth.Host = hostName
-	}
-	return nil
-}
-
 const minimalDbConnectVersion = "13.1"
 
 func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
@@ -208,3 +172,39 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 
 	return cmd
 }
+
+// Sets the host in the persistentAuth object based on the provided arguments and flags.
+// Follows the following precedence:
+// 1. [HOST] (first positional argument) or --host flag. Error if both are specified.
+// 2. Profile host, if available.
+// 3. Prompt the user for the host.
+func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
+	// If both [HOST] and --host are provided, return an error.
+	// TODO: test for error for this case.
+	if len(args) > 0 && persistentAuth.Host != "" {
+		return fmt.Errorf("please only provide a host as an argument or a flag, not both")
+	}
+
+	// If [HOST] is provided, set the host to the provided positional argument.
+	if len(args) > 0 && persistentAuth.Host == "" {
+		persistentAuth.Host = args[0]
+	}
+
+	// If neither [HOST] nor --host are provided, and the profile has a host, use it.
+	// Otherwise, prompt the user for a host.
+	if len(args) == 0 && persistentAuth.Host == "" {
+		hostName, err := getHostFromProfile(ctx, profileName)
+		if err != nil {
+			return err
+		}
+		if hostName == "" {
+			var err error
+			hostName, err = promptForHost(ctx)
+			if err != nil {
+				return err
+			}
+		}
+		persistentAuth.Host = hostName
+	}
+	return nil
+}

From daaf83c5c1e60d5393646a76b6824fc7cf1378d3 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:54:58 +0200
Subject: [PATCH 05/20] smaller diff

---
 cmd/auth/auth.go  | 34 ++++++++++++++++++++++++++++++++++
 cmd/auth/login.go | 27 ---------------------------
 2 files changed, 34 insertions(+), 27 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index c8e1118705..5feb328fe4 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -1,10 +1,44 @@
 package auth
 
 import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
 	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/spf13/cobra"
 )
 
+func promptForHost(ctx context.Context) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify a host using --host")
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks Host"
+	prompt.Validate = func(host string) error {
+		if !strings.HasPrefix(host, "https://") {
+			return fmt.Errorf("host URL must have a https:// prefix")
+		}
+		return nil
+	}
+	return prompt.Run()
+}
+
+func promptForAccountID(ctx context.Context) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		return "", errors.New("the command is being run in a non-interactive environment, please specify an account ID using --account-id")
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks Account ID"
+	prompt.Default = ""
+	prompt.AllowEdit = true
+	return prompt.Run()
+}
+
 func New() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "auth",
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index 55158ab3ae..d406817401 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -16,21 +16,6 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func promptForHost(ctx context.Context) (string, error) {
-	if !cmdio.IsInTTY(ctx) {
-		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify a host using --host")
-	}
-
-	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Host"
-	prompt.Validate = func(host string) error {
-		if !strings.HasPrefix(host, "https://") {
-			return fmt.Errorf("host URL must have a https:// prefix")
-		}
-		return nil
-	}
-	return prompt.Run()
-}
 
 func promptForProfile(ctx context.Context, dv string) (string, error) {
 	if !cmdio.IsInTTY(ctx) {
@@ -44,18 +29,6 @@ func promptForProfile(ctx context.Context, dv string) (string, error) {
 	return prompt.Run()
 }
 
-func promptForAccountID(ctx context.Context) (string, error) {
-	if !cmdio.IsInTTY(ctx) {
-		return "", errors.New("the command is being run in a non-interactive environment, please specify an account ID using --account-id")
-	}
-
-	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Account ID"
-	prompt.Default = ""
-	prompt.AllowEdit = true
-	return prompt.Run()
-}
-
 func getHostFromProfile(ctx context.Context, profileName string) (string, error) {
 	_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
 		return p.Name == profileName

From 580666981334986b4f15fb993fe10e86760bee4b Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:55:10 +0200
Subject: [PATCH 06/20] -

---
 cmd/auth/login.go | 2 --
 1 file changed, 2 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index d406817401..f1b394c559 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"strings"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
@@ -16,7 +15,6 @@ import (
 	"github.com/spf13/cobra"
 )
 
-
 func promptForProfile(ctx context.Context, dv string) (string, error) {
 	if !cmdio.IsInTTY(ctx) {
 		return "", errors.New("the command is being run in a non-interactive environment, please specify a profile using --profile")

From b91011ef1a9b35acb1dc4504e6099c4d65f94a97 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:55:45 +0200
Subject: [PATCH 07/20] -

---
 cmd/auth/auth.go | 50 ++++++++++++++++++++++++------------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 5feb328fe4..626b6604d3 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -11,6 +11,31 @@ import (
 	"github.com/spf13/cobra"
 )
 
+func New() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "auth",
+		Short: "Authentication related commands",
+		Long: `Authentication related commands. For more information regarding how
+authentication for the Databricks CLI and SDKs work please refer to the documentation
+linked below.
+
+AWS: https://docs.databricks.com/en/dev-tools/auth/index.html
+Azure: https://learn.microsoft.com/en-us/azure/databricks/dev-tools/auth
+GCP: https://docs.gcp.databricks.com/en/dev-tools/auth/index.html`,
+	}
+
+	var perisistentAuth auth.PersistentAuth
+	cmd.PersistentFlags().StringVar(&perisistentAuth.Host, "host", perisistentAuth.Host, "Databricks Host")
+	cmd.PersistentFlags().StringVar(&perisistentAuth.AccountID, "account-id", perisistentAuth.AccountID, "Databricks Account ID")
+
+	cmd.AddCommand(newEnvCommand())
+	cmd.AddCommand(newLoginCommand(&perisistentAuth))
+	cmd.AddCommand(newProfilesCommand())
+	cmd.AddCommand(newTokenCommand(&perisistentAuth))
+	cmd.AddCommand(newDescribeCommand())
+	return cmd
+}
+
 func promptForHost(ctx context.Context) (string, error) {
 	if !cmdio.IsInTTY(ctx) {
 		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify a host using --host")
@@ -38,28 +63,3 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.AllowEdit = true
 	return prompt.Run()
 }
-
-func New() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "auth",
-		Short: "Authentication related commands",
-		Long: `Authentication related commands. For more information regarding how
-authentication for the Databricks CLI and SDKs work please refer to the documentation
-linked below.
-
-AWS: https://docs.databricks.com/en/dev-tools/auth/index.html
-Azure: https://learn.microsoft.com/en-us/azure/databricks/dev-tools/auth
-GCP: https://docs.gcp.databricks.com/en/dev-tools/auth/index.html`,
-	}
-
-	var perisistentAuth auth.PersistentAuth
-	cmd.PersistentFlags().StringVar(&perisistentAuth.Host, "host", perisistentAuth.Host, "Databricks Host")
-	cmd.PersistentFlags().StringVar(&perisistentAuth.AccountID, "account-id", perisistentAuth.AccountID, "Databricks Account ID")
-
-	cmd.AddCommand(newEnvCommand())
-	cmd.AddCommand(newLoginCommand(&perisistentAuth))
-	cmd.AddCommand(newProfilesCommand())
-	cmd.AddCommand(newTokenCommand(&perisistentAuth))
-	cmd.AddCommand(newDescribeCommand())
-	return cmd
-}

From a6572dea971793811d82c2d4d9ce5477d0deb969 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 16 Apr 2024 18:57:13 +0200
Subject: [PATCH 08/20] -

---
 cmd/auth/login.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index f1b394c559..cb0a18079a 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -69,8 +69,7 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 
-		// If the user has not specified a profile name, and we are in an interactive
-		// session, prompt for one.
+		// If the user has not specified a profile name, prompt for one.
 		if profileName == "" {
 			var err error
 			profileName, err = promptForProfile(ctx, persistentAuth.DefaultProfileName())

From 1159315041c7623963c9003f249ceb6f05ecc6ce Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 17 Apr 2024 14:38:13 +0200
Subject: [PATCH 09/20] clean todos

---
 cmd/auth/login.go | 2 --
 1 file changed, 2 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index cb0a18079a..8416465802 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -34,7 +34,6 @@ func getHostFromProfile(ctx context.Context, profileName string) (string, error)
 
 	// Tolerate ErrNoConfiguration here, as we will write out a configuration file
 	// as part of the login flow.
-	// TODO: verify that a configuration file is written out as part of the login flow.
 	if err != nil && errors.Is(err, databrickscfg.ErrNoConfiguration) {
 		return "", nil
 	}
@@ -150,7 +149,6 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 // 3. Prompt the user for the host.
 func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
 	// If both [HOST] and --host are provided, return an error.
-	// TODO: test for error for this case.
 	if len(args) > 0 && persistentAuth.Host != "" {
 		return fmt.Errorf("please only provide a host as an argument or a flag, not both")
 	}

From 69ab91c4fcae2495a089161a33a1d53fda2c034c Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:10:13 +0200
Subject: [PATCH 10/20] remove errors.new

---
 cmd/auth/auth.go  | 2 +-
 cmd/auth/login.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 626b6604d3..f26bebf886 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -54,7 +54,7 @@ func promptForHost(ctx context.Context) (string, error) {
 
 func promptForAccountID(ctx context.Context) (string, error) {
 	if !cmdio.IsInTTY(ctx) {
-		return "", errors.New("the command is being run in a non-interactive environment, please specify an account ID using --account-id")
+		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify an account ID using --account-id")
 	}
 
 	prompt := cmdio.Prompt(ctx)
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index 8416465802..17f481fe75 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -17,7 +17,7 @@ import (
 
 func promptForProfile(ctx context.Context, dv string) (string, error) {
 	if !cmdio.IsInTTY(ctx) {
-		return "", errors.New("the command is being run in a non-interactive environment, please specify a profile using --profile")
+		return "", fmt.Errorf("the command is being run in a non-interactive environment, please specify a profile using --profile")
 	}
 
 	prompt := cmdio.Prompt(ctx)

From e31203ecc0e4e8fad2623adeb6945158f2aa9213 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:12:28 +0200
Subject: [PATCH 11/20] -

---
 cmd/auth/auth.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index f26bebf886..059fedd3f9 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -2,7 +2,6 @@ package auth
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"strings"
 

From 061957f3c28f4508c01ad8e09c4abd1753a3d6d6 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:16:52 +0200
Subject: [PATCH 12/20] remove profile flag override

---
 cmd/auth/login.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index 0a2f35d4e6..e7089a6b64 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -103,15 +103,14 @@ depends on the existing profiles you have set in your configuration file
 
 	var loginTimeout time.Duration
 	var configureCluster bool
-	var profileName string
 	cmd.Flags().DurationVar(&loginTimeout, "timeout", auth.DefaultTimeout,
 		"Timeout for completing login challenge in the browser")
 	cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
 		"Prompts to configure cluster")
-	cmd.Flags().StringVarP(&profileName, "profile", "p", "", `Name of the profile.`)
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
+		profileName := cmd.Flag("profile").Value.String()
 
 		// If the user has not specified a profile name, prompt for one.
 		if profileName == "" {

From fdaebad9659c3efcdfea5e57e9f9ad3aa026127c Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:22:26 +0200
Subject: [PATCH 13/20] revert pa.Close()

---
 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 e7089a6b64..a3083d23d9 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -126,6 +126,7 @@ depends on the existing profiles you have set in your configuration file
 		if err != nil {
 			return err
 		}
+		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.
@@ -149,7 +150,6 @@ depends on the existing profiles you have set in your configuration file
 		if err != nil {
 			return err
 		}
-		defer persistentAuth.Close()
 
 		if configureCluster {
 			w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))

From 74a77ff74cbe3be2613cae5763a9b7450c2d7911 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:34:29 +0200
Subject: [PATCH 14/20] add DEFAULT

---
 libs/auth/oauth.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go
index d75d0b683a..063f5fb9f4 100644
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@@ -95,6 +95,9 @@ func (a *PersistentAuth) DefaultProfileName() string {
 	if a.AccountID != "" {
 		return fmt.Sprintf("ACCOUNT-%s", a.AccountID)
 	}
+	if a.Host == "" {
+		return "DEFAULT"
+	}
 	host := strings.TrimPrefix(a.Host, "https://")
 	split := strings.Split(host, ".")
 	return split[0]

From 79f72bd472fcb93d7e77c90e475ea63c808d2741 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:42:45 +0200
Subject: [PATCH 15/20] validation for uuid

---
 cmd/auth/auth.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 17d81939b0..4ae22a87a5 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 )
 
@@ -60,5 +61,12 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.Label = "Databricks Account ID"
 	prompt.Default = ""
 	prompt.AllowEdit = true
+	prompt.Validate = func(accountID string) error {
+		_, err := uuid.Parse(accountID)
+		if err != nil {
+			return fmt.Errorf("account ID must be a valid UUID: %w", err)
+		}
+		return nil
+	}
 	return prompt.Run()
 }

From 2ccb79ca18de8bb234dffb905a20d77ddd343aa8 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Thu, 18 Apr 2024 15:46:39 +0200
Subject: [PATCH 16/20] lowercase

---
 cmd/auth/auth.go  | 4 ++--
 cmd/auth/login.go | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 4ae22a87a5..7d688832ff 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -42,7 +42,7 @@ func promptForHost(ctx context.Context) (string, error) {
 	}
 
 	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Host"
+	prompt.Label = "Databricks host"
 	prompt.Validate = func(host string) error {
 		if !strings.HasPrefix(host, "https://") {
 			return fmt.Errorf("host URL must have a https:// prefix")
@@ -58,7 +58,7 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	}
 
 	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Account ID"
+	prompt.Label = "Databricks account id"
 	prompt.Default = ""
 	prompt.AllowEdit = true
 	prompt.Validate = func(accountID string) error {
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index a3083d23d9..405c69efec 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -22,7 +22,7 @@ func promptForProfile(ctx context.Context, dv string) (string, error) {
 	}
 
 	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Profile Name"
+	prompt.Label = "Databricks profile name"
 	prompt.Default = dv
 	prompt.AllowEdit = true
 	return prompt.Run()

From 0766b9c97b327a752a31f631d57b719b6fa88b47 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 13 Aug 2024 17:57:01 +0200
Subject: [PATCH 17/20] fix tests

---
 cmd/auth/login_test.go | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go
index 4555f5800a..96a8550912 100644
--- a/cmd/auth/login_test.go
+++ b/cmd/auth/login_test.go
@@ -23,30 +23,30 @@ func TestSetHost(t *testing.T) {
 
 	// Test error when both flag and argument are provided
 	persistentAuth.Host = "val from --host"
-	err := setHost(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	err := setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
 	assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
 
 	// Test setting host from flag
 	persistentAuth.Host = "val from --host"
-	err = setHost(ctx, "profile-1", &persistentAuth, []string{})
+	err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "val from --host", persistentAuth.Host)
 
 	// Test setting host from argument
 	persistentAuth.Host = ""
-	err = setHost(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
 	assert.NoError(t, err)
 	assert.Equal(t, "val from [HOST]", persistentAuth.Host)
 
 	// Test setting host from profile
 	persistentAuth.Host = ""
-	err = setHost(ctx, "profile-1", &persistentAuth, []string{})
+	err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://www.host1.com", persistentAuth.Host)
 
 	// Test setting host from profile
 	persistentAuth.Host = ""
-	err = setHost(ctx, "profile-2", &persistentAuth, []string{})
+	err = setHostAndAccountId(ctx, "profile-2", &persistentAuth, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://www.host2.com", persistentAuth.Host)
 }

From d9f8288c1b1b501049b850d1b70e8a79cc8b403a Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 13 Aug 2024 18:12:14 +0200
Subject: [PATCH 18/20] add tests for account-id

---
 cmd/auth/login_test.go           | 35 +++++++++++++++++++++++++++++++-
 cmd/auth/testdata/.databrickscfg |  4 ++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go
index 96a8550912..d0fa5a16b8 100644
--- a/cmd/auth/login_test.go
+++ b/cmd/auth/login_test.go
@@ -5,8 +5,10 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
@@ -17,9 +19,9 @@ func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
 }
 
 func TestSetHost(t *testing.T) {
-	ctx := context.Background()
 	var persistentAuth auth.PersistentAuth
 	t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
+	ctx, _ := cmdio.SetupTest(context.Background())
 
 	// Test error when both flag and argument are provided
 	persistentAuth.Host = "val from --host"
@@ -49,4 +51,35 @@ func TestSetHost(t *testing.T) {
 	err = setHostAndAccountId(ctx, "profile-2", &persistentAuth, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://www.host2.com", persistentAuth.Host)
+
+	// Test host is not set. Should prompt.
+	persistentAuth.Host = ""
+	err = setHostAndAccountId(ctx, "", &persistentAuth, []string{})
+	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host")
+}
+
+func TestSetAccountId(t *testing.T) {
+	var persistentAuth auth.PersistentAuth
+	t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
+	ctx, _ := cmdio.SetupTest(context.Background())
+
+	// Test setting account-id from flag
+	persistentAuth.AccountID = "val from --account-id"
+	err := setHostAndAccountId(ctx, "account-profile", &persistentAuth, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
+	assert.Equal(t, "val from --account-id", persistentAuth.AccountID)
+
+	// Test setting account_id from profile
+	persistentAuth.AccountID = ""
+	err = setHostAndAccountId(ctx, "account-profile", &persistentAuth, []string{})
+	require.NoError(t, err)
+	assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
+	assert.Equal(t, "id-from-profile", persistentAuth.AccountID)
+
+	// Neither flag nor profile account-id is set, should prompt
+	persistentAuth.AccountID = ""
+	persistentAuth.Host = "https://accounts.cloud.databricks.com"
+	err = setHostAndAccountId(ctx, "", &persistentAuth, []string{})
+	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
 }
diff --git a/cmd/auth/testdata/.databrickscfg b/cmd/auth/testdata/.databrickscfg
index e815549bc3..06e55224a1 100644
--- a/cmd/auth/testdata/.databrickscfg
+++ b/cmd/auth/testdata/.databrickscfg
@@ -3,3 +3,7 @@ host = https://www.host1.com
 
 [profile-2]
 host = https://www.host2.com
+
+[account-profile]
+host = https://accounts.cloud.databricks.com
+account_id = id-from-profile

From 54bb1c33469b9b0335a7ec570d1021ca6f7fa774 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 13 Aug 2024 18:20:04 +0200
Subject: [PATCH 19/20] undo promptUI changes

---
 cmd/auth/auth.go   | 17 +----------------
 cmd/auth/login.go  |  2 +-
 libs/auth/oauth.go |  5 +----
 3 files changed, 3 insertions(+), 21 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 7d688832ff..128edd8411 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -3,11 +3,9 @@ package auth
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 )
 
@@ -42,13 +40,7 @@ func promptForHost(ctx context.Context) (string, error) {
 	}
 
 	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks host"
-	prompt.Validate = func(host string) error {
-		if !strings.HasPrefix(host, "https://") {
-			return fmt.Errorf("host URL must have a https:// prefix")
-		}
-		return nil
-	}
+	prompt.Label = "Databricks Host (e.g. https://<databricks-instance>.cloud.databricks.com)"
 	return prompt.Run()
 }
 
@@ -61,12 +53,5 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.Label = "Databricks account id"
 	prompt.Default = ""
 	prompt.AllowEdit = true
-	prompt.Validate = func(accountID string) error {
-		_, err := uuid.Parse(accountID)
-		if err != nil {
-			return fmt.Errorf("account ID must be a valid UUID: %w", err)
-		}
-		return nil
-	}
 	return prompt.Run()
 }
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index c1b2d924bd..00e357bfd8 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -112,7 +112,7 @@ depends on the existing profiles you have set in your configuration file
 		// If the user has not specified a profile name, prompt for one.
 		if profileName == "" {
 			var err error
-			profileName, err = promptForProfile(ctx, persistentAuth.DefaultProfileName())
+			profileName, err = promptForProfile(ctx, persistentAuth.ProfileName())
 			if err != nil {
 				return err
 			}
diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go
index ba8bb60fef..7c1cb95768 100644
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@@ -104,13 +104,10 @@ func (a *PersistentAuth) Load(ctx context.Context) (*oauth2.Token, error) {
 	return refreshed, nil
 }
 
-func (a *PersistentAuth) DefaultProfileName() string {
+func (a *PersistentAuth) ProfileName() string {
 	if a.AccountID != "" {
 		return fmt.Sprintf("ACCOUNT-%s", a.AccountID)
 	}
-	if a.Host == "" {
-		return "DEFAULT"
-	}
 	host := strings.TrimPrefix(a.Host, "https://")
 	split := strings.Split(host, ".")
 	return split[0]

From a475633aaa4c789405b7987de6fc10a3ad4f404c Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 13 Aug 2024 18:20:48 +0200
Subject: [PATCH 20/20] Revert "undo promptUI changes"

This reverts commit 54bb1c33469b9b0335a7ec570d1021ca6f7fa774.
---
 cmd/auth/auth.go   | 17 ++++++++++++++++-
 cmd/auth/login.go  |  2 +-
 libs/auth/oauth.go |  5 ++++-
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
index 128edd8411..7d688832ff 100644
--- a/cmd/auth/auth.go
+++ b/cmd/auth/auth.go
@@ -3,9 +3,11 @@ package auth
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 )
 
@@ -40,7 +42,13 @@ func promptForHost(ctx context.Context) (string, error) {
 	}
 
 	prompt := cmdio.Prompt(ctx)
-	prompt.Label = "Databricks Host (e.g. https://<databricks-instance>.cloud.databricks.com)"
+	prompt.Label = "Databricks host"
+	prompt.Validate = func(host string) error {
+		if !strings.HasPrefix(host, "https://") {
+			return fmt.Errorf("host URL must have a https:// prefix")
+		}
+		return nil
+	}
 	return prompt.Run()
 }
 
@@ -53,5 +61,12 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.Label = "Databricks account id"
 	prompt.Default = ""
 	prompt.AllowEdit = true
+	prompt.Validate = func(accountID string) error {
+		_, err := uuid.Parse(accountID)
+		if err != nil {
+			return fmt.Errorf("account ID must be a valid UUID: %w", err)
+		}
+		return nil
+	}
 	return prompt.Run()
 }
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index 00e357bfd8..c1b2d924bd 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -112,7 +112,7 @@ depends on the existing profiles you have set in your configuration file
 		// If the user has not specified a profile name, prompt for one.
 		if profileName == "" {
 			var err error
-			profileName, err = promptForProfile(ctx, persistentAuth.ProfileName())
+			profileName, err = promptForProfile(ctx, persistentAuth.DefaultProfileName())
 			if err != nil {
 				return err
 			}
diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go
index 7c1cb95768..ba8bb60fef 100644
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@@ -104,10 +104,13 @@ func (a *PersistentAuth) Load(ctx context.Context) (*oauth2.Token, error) {
 	return refreshed, nil
 }
 
-func (a *PersistentAuth) ProfileName() string {
+func (a *PersistentAuth) DefaultProfileName() string {
 	if a.AccountID != "" {
 		return fmt.Sprintf("ACCOUNT-%s", a.AccountID)
 	}
+	if a.Host == "" {
+		return "DEFAULT"
+	}
 	host := strings.TrimPrefix(a.Host, "https://")
 	split := strings.Split(host, ".")
 	return split[0]