[Feature] Migrate U2M Code to SDKs

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