mirror of https://github.com/databricks/cli.git
add DeleteKey to TokenCache for logout cmd
This commit is contained in:
parent
70ce802518
commit
882ccba0f5
|
@ -9,6 +9,7 @@ import (
|
||||||
type TokenCache interface {
|
type TokenCache interface {
|
||||||
Store(key string, t *oauth2.Token) error
|
Store(key string, t *oauth2.Token) error
|
||||||
Lookup(key string) (*oauth2.Token, error)
|
Lookup(key string) (*oauth2.Token, error)
|
||||||
|
DeleteKey(key string) error
|
||||||
}
|
}
|
||||||
|
|
||||||
var tokenCache int
|
var tokenCache int
|
||||||
|
|
|
@ -73,6 +73,29 @@ func (c *FileTokenCache) Lookup(key string) (*oauth2.Token, error) {
|
||||||
return t, nil
|
return t, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (c *FileTokenCache) DeleteKey(key string) error {
|
||||||
|
err := c.load()
|
||||||
|
if errors.Is(err, fs.ErrNotExist) {
|
||||||
|
return ErrNotConfigured
|
||||||
|
} else if err != nil {
|
||||||
|
return fmt.Errorf("load: %w", err)
|
||||||
|
}
|
||||||
|
c.Version = tokenCacheVersion
|
||||||
|
if c.Tokens == nil {
|
||||||
|
c.Tokens = map[string]*oauth2.Token{}
|
||||||
|
}
|
||||||
|
_, ok := c.Tokens[key]
|
||||||
|
if !ok {
|
||||||
|
return ErrNotConfigured
|
||||||
|
}
|
||||||
|
delete(c.Tokens, key)
|
||||||
|
raw, err := json.MarshalIndent(c, "", " ")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal: %w", err)
|
||||||
|
}
|
||||||
|
return os.WriteFile(c.fileLocation, raw, ownerReadWrite)
|
||||||
|
}
|
||||||
|
|
||||||
func (c *FileTokenCache) location() (string, error) {
|
func (c *FileTokenCache) location() (string, error) {
|
||||||
home, err := os.UserHomeDir()
|
home, err := os.UserHomeDir()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
|
@ -103,3 +103,29 @@ func TestStoreOnDev(t *testing.T) {
|
||||||
// macOS: read-only file system
|
// macOS: read-only file system
|
||||||
assert.Error(t, err)
|
assert.Error(t, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestStoreAndDeleteKey(t *testing.T) {
|
||||||
|
setup(t)
|
||||||
|
c := &FileTokenCache{}
|
||||||
|
err := c.Store("x", &oauth2.Token{
|
||||||
|
AccessToken: "abc",
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
err = c.Store("y", &oauth2.Token{
|
||||||
|
AccessToken: "bcd",
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
l := &FileTokenCache{}
|
||||||
|
err = l.DeleteKey("x")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 1, len(l.Tokens))
|
||||||
|
|
||||||
|
_, err = l.Lookup("x")
|
||||||
|
assert.Equal(t, ErrNotConfigured, err)
|
||||||
|
|
||||||
|
tok, err := l.Lookup("y")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "bcd", tok.AccessToken)
|
||||||
|
}
|
||||||
|
|
|
@ -23,4 +23,14 @@ func (i *InMemoryTokenCache) Store(key string, t *oauth2.Token) error {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// DeleteKey implements TokenCache.
|
||||||
|
func (i *InMemoryTokenCache) DeleteKey(key string) error {
|
||||||
|
_, ok := i.Tokens[key]
|
||||||
|
if !ok {
|
||||||
|
return ErrNotConfigured
|
||||||
|
}
|
||||||
|
delete(i.Tokens, key)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
var _ TokenCache = (*InMemoryTokenCache)(nil)
|
var _ TokenCache = (*InMemoryTokenCache)(nil)
|
||||||
|
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
"golang.org/x/oauth2"
|
"golang.org/x/oauth2"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -42,3 +43,29 @@ func TestInMemoryCacheStore(t *testing.T) {
|
||||||
assert.Equal(t, res, token)
|
assert.Equal(t, res, token)
|
||||||
assert.NoError(t, err)
|
assert.NoError(t, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestInMemoryDeleteKey(t *testing.T) {
|
||||||
|
c := &InMemoryTokenCache{
|
||||||
|
Tokens: map[string]*oauth2.Token{},
|
||||||
|
}
|
||||||
|
err := c.Store("x", &oauth2.Token{
|
||||||
|
AccessToken: "abc",
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
err = c.Store("y", &oauth2.Token{
|
||||||
|
AccessToken: "bcd",
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
err = c.DeleteKey("x")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 1, len(c.Tokens))
|
||||||
|
|
||||||
|
_, err = c.Lookup("x")
|
||||||
|
assert.Equal(t, ErrNotConfigured, err)
|
||||||
|
|
||||||
|
tok, err := c.Lookup("y")
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "bcd", tok.AccessToken)
|
||||||
|
}
|
||||||
|
|
|
@ -53,8 +53,9 @@ func TestOidcForWorkspace(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
type tokenCacheMock struct {
|
type tokenCacheMock struct {
|
||||||
store func(key string, t *oauth2.Token) error
|
store func(key string, t *oauth2.Token) error
|
||||||
lookup func(key string) (*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 {
|
func (m *tokenCacheMock) Store(key string, t *oauth2.Token) error {
|
||||||
|
@ -71,6 +72,13 @@ func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) {
|
||||||
return m.lookup(key)
|
return m.lookup(key)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (m *tokenCacheMock) DeleteKey(key string) error {
|
||||||
|
if m.deleteKey == nil {
|
||||||
|
panic("no deleteKey mock")
|
||||||
|
}
|
||||||
|
return m.deleteKey(key)
|
||||||
|
}
|
||||||
|
|
||||||
func TestLoad(t *testing.T) {
|
func TestLoad(t *testing.T) {
|
||||||
p := &PersistentAuth{
|
p := &PersistentAuth{
|
||||||
Host: "abc",
|
Host: "abc",
|
||||||
|
|
Loading…
Reference in New Issue