[Feature] Migrate U2M Code to SDKs (#2076)

## Changes
Ports the SDK to this change:
databricks/databricks-sdk-go#1108. This
eliminates the need for the SDK to invoke itself (!) when loading an
OAuth token. It also allows us to respond with more suitable errors when
fetching tokens. Functionally, this should behave identically to the
existing CLI from an external perspective.

One concrete improvement with this change is that the error message
returned when a refresh token is invalid now provides the correct
command to run, including the appropriate `--profile` flag.

Users will see the following message when using the CLI with a
profile/host with an expired refresh token:
![Screenshot 2025-01-20 at 13 18
59](https://github.com/user-attachments/assets/bf3040cb-971d-4c04-8557-0f6749d82b5e)


## Tests
Mostly this change involves refactoring existing code/tests to use/mock
the Go SDK interfaces, so there aren't too many substantial test
changes. `token_test.go` is rewritten using the table test approach to
reduce boilerplate.

Author: mgyucht

cmd/auth/auth.go

@@ -22,14 +22,14 @@ Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth
 GCP: https://docs.gcp.databricks.com/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")
+	var authArguments auth.AuthArguments
+	cmd.PersistentFlags().StringVar(&authArguments.Host, "host", "", "Databricks Host")
+	cmd.PersistentFlags().StringVar(&authArguments.AccountID, "account-id", "", "Databricks Account ID")
 
 	cmd.AddCommand(newEnvCommand())
-	cmd.AddCommand(newLoginCommand(&perisistentAuth))
+	cmd.AddCommand(newLoginCommand(&authArguments))
 	cmd.AddCommand(newProfilesCommand())
-	cmd.AddCommand(newTokenCommand(&perisistentAuth))
+	cmd.AddCommand(newTokenCommand(&authArguments))
 	cmd.AddCommand(newDescribeCommand())
 	return cmd
 }

cmd/auth/in_memory_test.go

@@ -0,0 +1,27 @@
+package auth
+
+import (
+	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"golang.org/x/oauth2"
+)
+
+type inMemoryTokenCache struct {
+	Tokens map[string]*oauth2.Token
+}
+
+// Lookup implements TokenCache.
+func (i *inMemoryTokenCache) Lookup(key string) (*oauth2.Token, error) {
+	token, ok := i.Tokens[key]
+	if !ok {
+		return nil, cache.ErrNotConfigured
+	}
+	return token, nil
+}
+
+// Store implements TokenCache.
+func (i *inMemoryTokenCache) Store(key string, t *oauth2.Token) error {
+	i.Tokens[key] = t
+	return nil
+}
+
+var _ cache.TokenCache = (*inMemoryTokenCache)(nil)

cmd/auth/login.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"runtime"
+	"strings"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
@@ -14,6 +15,7 @@ import (
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/spf13/cobra"
 )
 
@@ -34,7 +36,7 @@ const (
 	defaultTimeout          = 1 * time.Hour
 )
 
-func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
+func newLoginCommand(authArguments *auth.AuthArguments) *cobra.Command {
 	defaultConfigPath := "~/.databrickscfg"
 	if runtime.GOOS == "windows" {
 		defaultConfigPath = "%USERPROFILE%\\.databrickscfg"
@@ -98,14 +100,22 @@ 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, getProfileName(authArguments))
 			if err != nil {
 				return err
 			}
 		}
 
 		// Set the host and account-id based on the provided arguments and flags.
-		err := setHostAndAccountId(ctx, profileName, persistentAuth, args)
+		err := setHostAndAccountId(ctx, profile.DefaultProfiler, profileName, authArguments, args)
+		if err != nil {
+			return err
+		}
+		oauthArgument, err := authArguments.ToOAuthArgument()
+		if err != nil {
+			return err
+		}
+		persistentAuth, err := u2m.NewPersistentAuth(ctx, u2m.WithOAuthArgument(oauthArgument))
 		if err != nil {
 			return err
 		}
@@ -114,16 +124,15 @@ depends on the existing profiles you have set in your configuration file
 		// 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,
-			AccountID: persistentAuth.AccountID,
+			Host:      authArguments.Host,
+			AccountID: authArguments.AccountID,
 			AuthType:  "databricks-cli",
 		}
 
 		ctx, cancel := context.WithTimeout(ctx, loginTimeout)
 		defer cancel()
 
-		err = persistentAuth.Challenge(ctx)
-		if err != nil {
+		if err = persistentAuth.Challenge(); err != nil {
 			return err
 		}
 
@@ -173,53 +182,66 @@ depends on the existing profiles you have set in your configuration file
 // 1. --account-id flag.
 // 2. account-id from the specified profile, if available.
 // 3. Prompt the user for the account-id.
-func setHostAndAccountId(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
+func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profileName string, authArguments *auth.AuthArguments, args []string) error {
 	// If both [HOST] and --host are provided, return an error.
-	if len(args) > 0 && persistentAuth.Host != "" {
+	host := authArguments.Host
+	if len(args) > 0 && host != "" {
 		return errors.New("please only provide a host as an argument or a flag, not both")
 	}
 
-	profiler := profile.GetProfiler(ctx)
 	// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
 	profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName))
 	// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
 	if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
 		return err
 	}
 
-	if persistentAuth.Host == "" {
+	if host == "" {
 		if len(args) > 0 {
 			// If [HOST] is provided, set the host to the provided positional argument.
-			persistentAuth.Host = args[0]
+			authArguments.Host = args[0]
 		} else if len(profiles) > 0 && profiles[0].Host != "" {
 			// If neither [HOST] nor --host are provided, and the profile has a host, use it.
-			persistentAuth.Host = profiles[0].Host
+			authArguments.Host = profiles[0].Host
 		} else {
 			// If neither [HOST] nor --host are provided, and the profile does not have a host,
 			// then prompt the user for a host.
 			hostName, err := promptForHost(ctx)
 			if err != nil {
 				return err
 			}
-			persistentAuth.Host = hostName
+			authArguments.Host = hostName
 		}
 	}
 
 	// If the account-id was not provided as a cmd line flag, try to read it from
 	// the specified profile.
-	isAccountClient := (&config.Config{Host: persistentAuth.Host}).IsAccountClient()
-	if isAccountClient && persistentAuth.AccountID == "" {
+	isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient()
+	accountID := authArguments.AccountID
+	if isAccountClient && accountID == "" {
 		if len(profiles) > 0 && profiles[0].AccountID != "" {
-			persistentAuth.AccountID = profiles[0].AccountID
+			authArguments.AccountID = profiles[0].AccountID
 		} else {
 			// Prompt user for the account-id if it we could not get it from a
 			// profile.
 			accountId, err := promptForAccountID(ctx)
 			if err != nil {
 				return err
 			}
-			persistentAuth.AccountID = accountId
+			authArguments.AccountID = accountId
 		}
 	}
 	return nil
 }
+
+// getProfileName returns the default profile name for a given host/account ID.
+// If the account ID is provided, the profile name is "ACCOUNT-<account-id>".
+// Otherwise, the profile name is the first part of the host URL.
+func getProfileName(authArguments *auth.AuthArguments) string {
+	if authArguments.AccountID != "" {
+		return "ACCOUNT-" + authArguments.AccountID
+	}
+	host := strings.TrimPrefix(authArguments.Host, "https://")
+	split := strings.Split(host, ".")
+	return split[0]
+}

cmd/auth/login_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -14,72 +15,72 @@ import (
 func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
 	ctx := context.Background()
 	ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg")
-	err := setHostAndAccountId(ctx, "foo", &auth.PersistentAuth{Host: "test"}, []string{})
+	err := setHostAndAccountId(ctx, profile.DefaultProfiler, "foo", &auth.AuthArguments{Host: "test"}, []string{})
 	assert.NoError(t, err)
 }
 
 func TestSetHost(t *testing.T) {
-	var persistentAuth auth.PersistentAuth
+	authArguments := auth.AuthArguments{}
 	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"
-	err := setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	authArguments.Host = "val from --host"
+	err := setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []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 = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
+	authArguments.Host = "val from --host"
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
 	assert.NoError(t, err)
-	assert.Equal(t, "val from --host", persistentAuth.Host)
+	assert.Equal(t, "val from --host", authArguments.Host)
 
 	// Test setting host from argument
-	persistentAuth.Host = ""
-	err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
+	authArguments.Host = ""
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"})
 	assert.NoError(t, err)
-	assert.Equal(t, "val from [HOST]", persistentAuth.Host)
+	assert.Equal(t, "val from [HOST]", authArguments.Host)
 
 	// Test setting host from profile
-	persistentAuth.Host = ""
-	err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
+	authArguments.Host = ""
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
 	assert.NoError(t, err)
-	assert.Equal(t, "https://www.host1.com", persistentAuth.Host)
+	assert.Equal(t, "https://www.host1.com", authArguments.Host)
 
 	// Test setting host from profile
-	persistentAuth.Host = ""
-	err = setHostAndAccountId(ctx, "profile-2", &persistentAuth, []string{})
+	authArguments.Host = ""
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-2", &authArguments, []string{})
 	assert.NoError(t, err)
-	assert.Equal(t, "https://www.host2.com", persistentAuth.Host)
+	assert.Equal(t, "https://www.host2.com", authArguments.Host)
 
 	// Test host is not set. Should prompt.
-	persistentAuth.Host = ""
-	err = setHostAndAccountId(ctx, "", &persistentAuth, []string{})
+	authArguments.Host = ""
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []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
+	var authArguments auth.AuthArguments
 	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{})
+	authArguments.AccountID = "val from --account-id"
+	err := setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
 	assert.NoError(t, err)
-	assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
-	assert.Equal(t, "val from --account-id", persistentAuth.AccountID)
+	assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
+	assert.Equal(t, "val from --account-id", authArguments.AccountID)
 
 	// Test setting account_id from profile
-	persistentAuth.AccountID = ""
-	err = setHostAndAccountId(ctx, "account-profile", &persistentAuth, []string{})
+	authArguments.AccountID = ""
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
 	require.NoError(t, err)
-	assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
-	assert.Equal(t, "id-from-profile", persistentAuth.AccountID)
+	assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
+	assert.Equal(t, "id-from-profile", authArguments.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{})
+	authArguments.AccountID = ""
+	authArguments.Host = "https://accounts.cloud.databricks.com"
+	err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{})
 	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
 }