use PersistentAuth struc

cmd/auth/auth.go

@@ -31,7 +31,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`,
 	cmd.AddCommand(newProfilesCommand())
 	cmd.AddCommand(newTokenCommand(&perisistentAuth))
 	cmd.AddCommand(newDescribeCommand())
-	cmd.AddCommand(newLogoutCommand())
+	cmd.AddCommand(newLogoutCommand(&perisistentAuth))
 	return cmd
 }
 

cmd/auth/logout.go

@@ -6,7 +6,7 @@ import (
 	"fmt"
 	"io/fs"
 
-	"github.com/databricks/cli/libs/auth/cache"
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
@@ -14,15 +14,15 @@ import (
 	"github.com/spf13/cobra"
 )
 
-type Logout struct {
+type LogoutSession struct {
 	Profile        string
 	File           config.File
-	Cache   cache.TokenCache
+	PersistentAuth *auth.PersistentAuth
 }
 
-func (l *Logout) load(ctx context.Context, profileName string) error {
+func (l *LogoutSession) load(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth) error {
 	l.Profile = profileName
-	l.Cache = cache.GetTokenCache(ctx)
+	l.PersistentAuth = persistentAuth
 	iniFile, err := profile.DefaultProfiler.Get(ctx)
 	if errors.Is(err, fs.ErrNotExist) {
 		return err
@@ -33,7 +33,7 @@ func (l *Logout) load(ctx context.Context, profileName string) error {
 	return nil
 }
 
-func (l *Logout) getSetionMap() (map[string]string, error) {
+func (l *LogoutSession) getConfigSetionMap() (map[string]string, error) {
 	section, err := l.File.GetSection(l.Profile)
 	if err != nil {
 		return map[string]string{}, fmt.Errorf("profile does not exist in config file: %w", err)
@@ -42,13 +42,13 @@ func (l *Logout) getSetionMap() (map[string]string, error) {
 }
 
 // clear token from ~/.databricks/token-cache.json
-func (l *Logout) clearTokenCache(key string) error {
-	return l.Cache.DeleteKey(key)
+func (l *LogoutSession) clearTokenCache(ctx context.Context) error {
+	return l.PersistentAuth.ClearToken(ctx)
 }
 
 // Overrewrite profile to .databrickscfg without fields marked as sensitive
 // Other attributes are preserved.
-func (l *Logout) clearConfigFile(ctx context.Context, sectionMap map[string]string) error {
+func (l *LogoutSession) clearConfigFile(ctx context.Context, sectionMap map[string]string) error {
 	return databrickscfg.SaveToProfile(ctx, &config.Config{
 		ConfigFile:           l.File.Path(),
 		Profile:              l.Profile,
@@ -69,7 +69,7 @@ func (l *Logout) clearConfigFile(ctx context.Context, sectionMap map[string]stri
 	})
 }
 
-func newLogoutCommand() *cobra.Command {
+func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "logout [PROFILE]",
 		Short: "Logout from specified profile",
@@ -77,22 +77,31 @@ func newLogoutCommand() *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		var profileName string
-		if len(args) < 1 {
-			profileName = cmd.Flag("profile").Value.String()
-		} else {
-			profileName = args[0]
-		}
-		logout := &Logout{}
-		logout.load(ctx, profileName)
-		sectionMap, err := logout.getSetionMap()
+		profileName := cmd.Flag("profile").Value.String()
+		// If the user has not specified a profile name, prompt for one.
+		if profileName == "" {
+			var err error
+			profileName, err = promptForProfile(ctx, persistentAuth.ProfileName())
 			if err != nil {
 				return err
 			}
-		if err := logout.clearTokenCache(sectionMap["host"]); err != nil {
+		}
+		// Set the host and account-id based on the provided arguments and flags.
+		err := setHostAndAccountId(ctx, profileName, persistentAuth, args)
+		if err != nil {
 			return err
 		}
-		if err := logout.clearConfigFile(ctx, sectionMap); err != nil {
+		defer persistentAuth.Close()
+		LogoutSession := &LogoutSession{}
+		LogoutSession.load(ctx, profileName, persistentAuth)
+		configSectionMap, err := LogoutSession.getConfigSetionMap()
+		if err != nil {
+			return err
+		}
+		if err := LogoutSession.clearTokenCache(ctx); err != nil {
+			return err
+		}
+		if err := LogoutSession.clearConfigFile(ctx, configSectionMap); err != nil {
 			return err
 		}
 		cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully logged out", profileName))

cmd/auth/logout_test.go

@@ -2,13 +2,11 @@ package auth
 
 import (
 	"context"
-	"fmt"
 	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/oauth2"
 
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/databricks-sdk-go/config"
@@ -27,7 +25,7 @@ func TestLogout_ClearConfigFile(t *testing.T) {
 	require.NoError(t, err)
 	iniFile, err := config.LoadFile(path)
 	require.NoError(t, err)
-	logout := &Logout{
+	logout := &LogoutSession{
 		Profile: "abc",
 		File:    *iniFile,
 	}
@@ -49,65 +47,3 @@ func TestLogout_ClearConfigFile(t *testing.T) {
 	assert.Len(t, raw, 1)
 	assert.Equal(t, "https://foo", raw["host"])
 }
-
-type tokenCacheMock struct {
-	store     func(key string, t *oauth2.Token) error
-	lookup    func(key string) (*oauth2.Token, error)
-	deleteKey func(key string) error
-}
-
-func (m *tokenCacheMock) Store(key string, t *oauth2.Token) error {
-	if m.store == nil {
-		panic("no store mock")
-	}
-	return m.store(key, t)
-}
-
-func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) {
-	if m.lookup == nil {
-		panic("no lookup mock")
-	}
-	return m.lookup(key)
-}
-
-func (m *tokenCacheMock) DeleteKey(key string) error {
-	if m.deleteKey == nil {
-		panic("no deleteKey mock")
-	}
-	return m.deleteKey(key)
-}
-
-func TestLogout_ClearTokenCache(t *testing.T) {
-	ctx := context.Background()
-	path := filepath.Join(t.TempDir(), "databrickscfg")
-
-	err := databrickscfg.SaveToProfile(ctx, &config.Config{
-		ConfigFile: path,
-		Profile:    "abc",
-		Host:       "https://foo",
-		AuthType:   "databricks-cli",
-	})
-	require.NoError(t, err)
-	iniFile, err := config.LoadFile(path)
-	require.NoError(t, err)
-	logout := &Logout{
-		Profile: "abc",
-		File:    *iniFile,
-		Cache: &tokenCacheMock{
-			deleteKey: func(key string) error {
-				assert.Equal(t, "https://foo", key)
-				return nil
-			},
-			lookup: func(key string) (*oauth2.Token, error) {
-				assert.Equal(t, "https://foo", key)
-				return &oauth2.Token{}, fmt.Errorf("No token found")
-			},
-		},
-	}
-	sectionMap, err := logout.getSetionMap()
-	assert.NoError(t, err)
-	err = logout.clearTokenCache(sectionMap["host"])
-	assert.NoError(t, err)
-	_, err = logout.Cache.Lookup(sectionMap["host"])
-	assert.Error(t, err)
-}

libs/auth/cache/file_test.go

@@ -129,3 +129,14 @@ func TestStoreAndDeleteKey(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "bcd", tok.AccessToken)
 }
+
+func TestDeleteKeyNotExist(t *testing.T) {
+	c := &FileTokenCache{
+		Tokens: map[string]*oauth2.Token{},
+	}
+	err := c.DeleteKey("x")
+	assert.Equal(t, ErrNotConfigured, err)
+
+	_, err = c.Lookup("x")
+	assert.Equal(t, ErrNotConfigured, err)
+}

libs/auth/cache/in_memory_test.go

@@ -69,3 +69,14 @@ func TestInMemoryDeleteKey(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "bcd", tok.AccessToken)
 }
+
+func TestInMemoryDeleteKeyNotExist(t *testing.T) {
+	c := &InMemoryTokenCache{
+		Tokens: map[string]*oauth2.Token{},
+	}
+    err := c.DeleteKey("x")
+	assert.Equal(t, ErrNotConfigured, err)
+
+	_, err = c.Lookup("x")
+	assert.Equal(t, ErrNotConfigured, err)
+}

libs/auth/oauth.go

@@ -143,6 +143,18 @@ func (a *PersistentAuth) Challenge(ctx context.Context) error {
 	return nil
 }
 
+func (a *PersistentAuth) ClearToken(ctx context.Context) error {
+	if a.Host == "" && a.AccountID == "" {
+		return ErrFetchCredentials
+	}
+	if a.cache == nil {
+		a.cache = cache.GetTokenCache(ctx)
+	}
+	// lookup token identified by host (and possibly the account id)
+	key := a.key()
+	return a.cache.DeleteKey(key)
+}
+
 func (a *PersistentAuth) init(ctx context.Context) error {
 	if a.Host == "" && a.AccountID == "" {
 		return ErrFetchCredentials

libs/auth/oauth_test.go

@@ -236,3 +236,49 @@ func TestChallengeFailed(t *testing.T) {
 		assert.EqualError(t, err, "authorize: access_denied: Policy evaluation failed for this request")
 	})
 }
+
+func TestClearToken(t *testing.T) {
+	p := &PersistentAuth{
+		Host:      "abc",
+		AccountID: "xyz",
+		cache: &tokenCacheMock{
+			lookup: func(key string) (*oauth2.Token, error) {
+				assert.Equal(t, "https://abc/oidc/accounts/xyz", key)
+				return &oauth2.Token{}, ErrNotConfigured
+			},
+			deleteKey: func(key string) error {
+				assert.Equal(t, "https://abc/oidc/accounts/xyz", key)
+				return nil
+			},
+		},
+	}
+	defer p.Close()
+	err := p.ClearToken(context.Background())
+	assert.NoError(t, err)
+	key := p.key()
+	_, err = p.cache.Lookup(key)
+	assert.Equal(t, ErrNotConfigured, err)
+}
+
+func TestClearTokenNotExist(t *testing.T) {
+	p := &PersistentAuth{
+		Host:      "abc",
+		AccountID: "xyz",
+		cache: &tokenCacheMock{
+			lookup: func(key string) (*oauth2.Token, error) {
+				assert.Equal(t, "https://abc/oidc/accounts/xyz", key)
+				return &oauth2.Token{}, ErrNotConfigured
+			},
+			deleteKey: func(key string) error {
+				assert.Equal(t, "https://abc/oidc/accounts/xyz", key)
+				return ErrNotConfigured
+			},
+		},
+	}
+	defer p.Close()
+	err := p.ClearToken(context.Background())
+	assert.Equal(t, ErrNotConfigured, err)
+	key := p.key()
+	_, err = p.cache.Lookup(key)
+	assert.Equal(t, ErrNotConfigured, err)
+}