use PersistentAuth struc

This commit is contained in:
Richard Nordström 2024-09-01 18:22:18 +02:00
parent 712e2919f5
commit 6a8b2f452f
No known key found for this signature in database
GPG Key ID: ACCB352EC60AF27C
7 changed files with 113 additions and 88 deletions

View File

@ -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
}

View File

@ -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 {
Profile string
File config.File
Cache cache.TokenCache
type LogoutSession struct {
Profile string
File config.File
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]
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
}
}
logout := &Logout{}
logout.load(ctx, profileName)
sectionMap, err := logout.getSetionMap()
// 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.clearTokenCache(sectionMap["host"]); err != nil {
defer persistentAuth.Close()
LogoutSession := &LogoutSession{}
LogoutSession.load(ctx, profileName, persistentAuth)
configSectionMap, err := LogoutSession.getConfigSetionMap()
if err != nil {
return err
}
if err := logout.clearConfigFile(ctx, sectionMap); err != nil {
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))

View File

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

View File

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

View File

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

View File

@ -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

View File

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