From 882ccba0f59c8c47688843fc53945030b5cebd3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sat, 24 Aug 2024 23:32:30 +0200 Subject: [PATCH 01/13] add DeleteKey to TokenCache for logout cmd --- libs/auth/cache/cache.go | 1 + libs/auth/cache/file.go | 23 +++++++++++++++++++++++ libs/auth/cache/file_test.go | 26 ++++++++++++++++++++++++++ libs/auth/cache/in_memory.go | 10 ++++++++++ libs/auth/cache/in_memory_test.go | 27 +++++++++++++++++++++++++++ libs/auth/oauth_test.go | 12 ++++++++++-- 6 files changed, 97 insertions(+), 2 deletions(-) diff --git a/libs/auth/cache/cache.go b/libs/auth/cache/cache.go index 097353e7..7ca9e8b6 100644 --- a/libs/auth/cache/cache.go +++ b/libs/auth/cache/cache.go @@ -9,6 +9,7 @@ import ( type TokenCache interface { Store(key string, t *oauth2.Token) error Lookup(key string) (*oauth2.Token, error) + DeleteKey(key string) error } var tokenCache int diff --git a/libs/auth/cache/file.go b/libs/auth/cache/file.go index 38dfea9f..9e99070e 100644 --- a/libs/auth/cache/file.go +++ b/libs/auth/cache/file.go @@ -73,6 +73,29 @@ func (c *FileTokenCache) Lookup(key string) (*oauth2.Token, error) { 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) { home, err := os.UserHomeDir() if err != nil { diff --git a/libs/auth/cache/file_test.go b/libs/auth/cache/file_test.go index 3e4aae36..3d0801a5 100644 --- a/libs/auth/cache/file_test.go +++ b/libs/auth/cache/file_test.go @@ -103,3 +103,29 @@ func TestStoreOnDev(t *testing.T) { // macOS: read-only file system 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) +} diff --git a/libs/auth/cache/in_memory.go b/libs/auth/cache/in_memory.go index 469d4557..6daaf868 100644 --- a/libs/auth/cache/in_memory.go +++ b/libs/auth/cache/in_memory.go @@ -23,4 +23,14 @@ func (i *InMemoryTokenCache) Store(key string, t *oauth2.Token) error { 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) diff --git a/libs/auth/cache/in_memory_test.go b/libs/auth/cache/in_memory_test.go index d8394d3b..f67bf8d3 100644 --- a/libs/auth/cache/in_memory_test.go +++ b/libs/auth/cache/in_memory_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) @@ -42,3 +43,29 @@ func TestInMemoryCacheStore(t *testing.T) { assert.Equal(t, res, token) 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) +} diff --git a/libs/auth/oauth_test.go b/libs/auth/oauth_test.go index ea6a8061..42cc0ef9 100644 --- a/libs/auth/oauth_test.go +++ b/libs/auth/oauth_test.go @@ -53,8 +53,9 @@ func TestOidcForWorkspace(t *testing.T) { } type tokenCacheMock struct { - store func(key string, t *oauth2.Token) error - lookup func(key string) (*oauth2.Token, error) + 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 { @@ -71,6 +72,13 @@ func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) { 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) { p := &PersistentAuth{ Host: "abc", From 712e2919f5d6d1b7df396bb067a2e3c853d0408d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 25 Aug 2024 23:53:44 +0200 Subject: [PATCH 02/13] add logout cmd --- cmd/auth/auth.go | 1 + cmd/auth/logout.go | 102 ++++++++++++++++++++++++++++++++++++ cmd/auth/logout_test.go | 113 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 216 insertions(+) create mode 100644 cmd/auth/logout.go create mode 100644 cmd/auth/logout_test.go diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index ceceae25..57b7b8a4 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -31,6 +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()) return cmd } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go new file mode 100644 index 00000000..1e51645e --- /dev/null +++ b/cmd/auth/logout.go @@ -0,0 +1,102 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "io/fs" + + "github.com/databricks/cli/libs/auth/cache" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" +) + +type Logout struct { + Profile string + File config.File + Cache cache.TokenCache +} + +func (l *Logout) load(ctx context.Context, profileName string) error { + l.Profile = profileName + l.Cache = cache.GetTokenCache(ctx) + iniFile, err := profile.DefaultProfiler.Get(ctx) + if errors.Is(err, fs.ErrNotExist) { + return err + } else if err != nil { + return fmt.Errorf("cannot parse config file: %w", err) + } + l.File = *iniFile + return nil +} + +func (l *Logout) getSetionMap() (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) + } + return section.KeysHash(), nil +} + +// clear token from ~/.databricks/token-cache.json +func (l *Logout) clearTokenCache(key string) error { + return l.Cache.DeleteKey(key) +} + +// 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 { + return databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: l.File.Path(), + Profile: l.Profile, + Host: sectionMap["host"], + ClusterID: sectionMap["cluster_id"], + WarehouseID: sectionMap["warehouse_id"], + ServerlessComputeID: sectionMap["serverless_compute_id"], + AccountID: sectionMap["account_id"], + Username: sectionMap["username"], + GoogleServiceAccount: sectionMap["google_service_account"], + AzureResourceID: sectionMap["azure_workspace_resource_id"], + AzureClientID: sectionMap["azure_client_id"], + AzureTenantID: sectionMap["azure_tenant_id"], + AzureEnvironment: sectionMap["azure_environment"], + AzureLoginAppID: sectionMap["azure_login_app_id"], + ClientID: sectionMap["client_id"], + AuthType: sectionMap["auth_type"], + }) +} + +func newLogoutCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "logout [PROFILE]", + Short: "Logout from specified profile", + } + + 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() + if err != nil { + return err + } + if err := logout.clearTokenCache(sectionMap["host"]); err != nil { + return err + } + if err := logout.clearConfigFile(ctx, sectionMap); err != nil { + return err + } + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully logged out", profileName)) + return nil + } + return cmd +} diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go new file mode 100644 index 00000000..e3344489 --- /dev/null +++ b/cmd/auth/logout_test.go @@ -0,0 +1,113 @@ +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" +) + +func TestLogout_ClearConfigFile(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", + Token: "xyz", + }) + require.NoError(t, err) + iniFile, err := config.LoadFile(path) + require.NoError(t, err) + logout := &Logout{ + Profile: "abc", + File: *iniFile, + } + section, err := logout.File.GetSection("abc") + assert.NoError(t, err) + sectionMap := section.KeysHash() + err = logout.clearConfigFile(ctx, sectionMap) + assert.NoError(t, err) + + iniFile, err = config.LoadFile(path) + require.NoError(t, err) + assert.Len(t, iniFile.Sections(), 2) + assert.True(t, iniFile.HasSection("DEFAULT")) + assert.True(t, iniFile.HasSection("abc")) + + abc, err := iniFile.GetSection("abc") + assert.NoError(t, err) + raw := abc.KeysHash() + 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) +} From 6a8b2f452f7c536257b187bbd2eca3e99a7a2fd3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 1 Sep 2024 18:22:18 +0200 Subject: [PATCH 03/13] use PersistentAuth struc --- cmd/auth/auth.go | 2 +- cmd/auth/logout.go | 53 ++++++++++++++----------- cmd/auth/logout_test.go | 66 +------------------------------ libs/auth/cache/file_test.go | 11 ++++++ libs/auth/cache/in_memory_test.go | 11 ++++++ libs/auth/oauth.go | 12 ++++++ libs/auth/oauth_test.go | 46 +++++++++++++++++++++ 7 files changed, 113 insertions(+), 88 deletions(-) diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 57b7b8a4..967e79a2 100644 --- a/cmd/auth/auth.go +++ b/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 } diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 1e51645e..586d37ef 100644 --- a/cmd/auth/logout.go +++ b/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 { - 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)) diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index e3344489..73f99f2a 100644 --- a/cmd/auth/logout_test.go +++ b/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) -} diff --git a/libs/auth/cache/file_test.go b/libs/auth/cache/file_test.go index 3d0801a5..fb7375c2 100644 --- a/libs/auth/cache/file_test.go +++ b/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) +} diff --git a/libs/auth/cache/in_memory_test.go b/libs/auth/cache/in_memory_test.go index f67bf8d3..3714ed79 100644 --- a/libs/auth/cache/in_memory_test.go +++ b/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) +} diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go index 7c1cb957..46f2af11 100644 --- a/libs/auth/oauth.go +++ b/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 diff --git a/libs/auth/oauth_test.go b/libs/auth/oauth_test.go index 42cc0ef9..e5e36032 100644 --- a/libs/auth/oauth_test.go +++ b/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) +} From 6277cf24c6558e72a29527651b05a826f699200b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 1 Sep 2024 20:17:07 +0200 Subject: [PATCH 04/13] allow no OAuth in case PAT is used --- cmd/auth/logout.go | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 586d37ef..506c97ee 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -6,6 +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" @@ -98,8 +99,14 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { if err != nil { return err } - if err := LogoutSession.clearTokenCache(ctx); err != nil { - return err + err = LogoutSession.clearTokenCache(ctx) + if err != nil { + if errors.Is(err, cache.ErrNotConfigured) { + // It is OK to not have OAuth configured. Move on and remove + // sensitive values example PAT from config file + } else { + return err + } } if err := LogoutSession.clearConfigFile(ctx, configSectionMap); err != nil { return err From 7eca34a7b25f835ac184cd4784aba07517bcfc16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 1 Sep 2024 20:25:40 +0200 Subject: [PATCH 05/13] fix typo --- cmd/auth/logout.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 506c97ee..05a1843a 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -34,7 +34,7 @@ func (l *LogoutSession) load(ctx context.Context, profileName string, persistent return nil } -func (l *LogoutSession) getConfigSetionMap() (map[string]string, error) { +func (l *LogoutSession) getConfigSectionMap() (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) @@ -95,7 +95,7 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { defer persistentAuth.Close() LogoutSession := &LogoutSession{} LogoutSession.load(ctx, profileName, persistentAuth) - configSectionMap, err := LogoutSession.getConfigSetionMap() + configSectionMap, err := LogoutSession.getConfigSectionMap() if err != nil { return err } From 6c32a0df7abe901f0bf707d825aa027a9a40e62b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 2 Sep 2024 00:15:48 +0200 Subject: [PATCH 06/13] improve profile handling and add tests --- cmd/auth/logout.go | 47 ++++++++++++++++++++++++++++----------- cmd/auth/logout_test.go | 49 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 13 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 05a1843a..37fba16f 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -6,8 +6,8 @@ import ( "fmt" "io/fs" - "github.com/databricks/cli/libs/auth/cache" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/cache" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" @@ -31,6 +31,22 @@ func (l *LogoutSession) load(ctx context.Context, profileName string, persistent return fmt.Errorf("cannot parse config file: %w", err) } l.File = *iniFile + if err := l.setHostAndAccountIdFromProfile(); err != nil { + return err + } + return nil +} + +func (l *LogoutSession) setHostAndAccountIdFromProfile() error { + sectionMap, err := l.getConfigSectionMap() + if err != nil { + return err + } + if sectionMap["host"] == "" { + return fmt.Errorf("no host configured for profile %s", l.Profile) + } + l.PersistentAuth.Host = sectionMap["host"] + l.PersistentAuth.AccountID = sectionMap["account_id"] return nil } @@ -74,11 +90,21 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { cmd := &cobra.Command{ Use: "logout [PROFILE]", Short: "Logout from specified profile", + Long: "Clears OAuth token from token-cache and any sensitive value in the config file, if they exist.", } cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - profileName := cmd.Flag("profile").Value.String() + profileNameFromFlag := cmd.Flag("profile").Value.String() + // If both [PROFILE] and --profile are provided, return an error. + if len(args) > 0 && profileNameFromFlag != "" { + return fmt.Errorf("please only provide a profile as an argument or a flag, not both") + } + // Determine the profile name from either args or the flag. + profileName := profileNameFromFlag + if len(args) > 0 { + profileName = args[0] + } // If the user has not specified a profile name, prompt for one. if profileName == "" { var err error @@ -87,28 +113,23 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { return err } } - // Set the host and account-id based on the provided arguments and flags. - err := setHostAndAccountId(ctx, profileName, persistentAuth, args) - if err != nil { - return err - } defer persistentAuth.Close() - LogoutSession := &LogoutSession{} - LogoutSession.load(ctx, profileName, persistentAuth) - configSectionMap, err := LogoutSession.getConfigSectionMap() + logoutSession := &LogoutSession{} + logoutSession.load(ctx, profileName, persistentAuth) + configSectionMap, err := logoutSession.getConfigSectionMap() if err != nil { return err } - err = LogoutSession.clearTokenCache(ctx) + err = logoutSession.clearTokenCache(ctx) if err != nil { if errors.Is(err, cache.ErrNotConfigured) { // It is OK to not have OAuth configured. Move on and remove - // sensitive values example PAT from config file + // sensitive values from config file (Example PAT) } else { return err } } - if err := LogoutSession.clearConfigFile(ctx, configSectionMap); err != nil { + if err := logoutSession.clearConfigFile(ctx, configSectionMap); err != nil { return err } cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully logged out", profileName)) diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index 73f99f2a..4f1133a3 100644 --- a/cmd/auth/logout_test.go +++ b/cmd/auth/logout_test.go @@ -8,6 +8,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/databricks-sdk-go/config" ) @@ -47,3 +48,51 @@ func TestLogout_ClearConfigFile(t *testing.T) { assert.Len(t, raw, 1) assert.Equal(t, "https://foo", raw["host"]) } + +func TestLogout_setHostAndAccountIdFromProfile(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", + Token: "xyz", + }) + require.NoError(t, err) + iniFile, err := config.LoadFile(path) + require.NoError(t, err) + logout := &LogoutSession{ + Profile: "abc", + File: *iniFile, + PersistentAuth: &auth.PersistentAuth{}, + } + err = logout.setHostAndAccountIdFromProfile() + assert.NoError(t, err) + assert.Equal(t, logout.PersistentAuth.Host, "https://foo") + assert.Empty(t, logout.PersistentAuth.AccountID) +} + +func TestLogout_getConfigSectionMap(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", + Token: "xyz", + }) + require.NoError(t, err) + iniFile, err := config.LoadFile(path) + require.NoError(t, err) + logout := &LogoutSession{ + Profile: "abc", + File: *iniFile, + PersistentAuth: &auth.PersistentAuth{}, + } + configSectionMap, err := logout.getConfigSectionMap() + assert.NoError(t, err) + assert.Equal(t, configSectionMap["host"], "https://foo") + assert.Equal(t, configSectionMap["token"], "xyz") +} From 37067ef93391e4e8da6b5b71f7d2be9ad1622883 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 20:21:38 +0200 Subject: [PATCH 07/13] rename DeleteKey to Delete --- libs/auth/cache/cache.go | 2 +- libs/auth/cache/file.go | 2 +- libs/auth/cache/file_test.go | 4 ++-- libs/auth/cache/in_memory.go | 4 ++-- libs/auth/cache/in_memory_test.go | 4 ++-- libs/auth/oauth.go | 2 +- libs/auth/oauth_test.go | 2 +- 7 files changed, 10 insertions(+), 10 deletions(-) diff --git a/libs/auth/cache/cache.go b/libs/auth/cache/cache.go index 7ca9e8b6..2c88d093 100644 --- a/libs/auth/cache/cache.go +++ b/libs/auth/cache/cache.go @@ -9,7 +9,7 @@ import ( type TokenCache interface { Store(key string, t *oauth2.Token) error Lookup(key string) (*oauth2.Token, error) - DeleteKey(key string) error + Delete(key string) error } var tokenCache int diff --git a/libs/auth/cache/file.go b/libs/auth/cache/file.go index 9e99070e..4a7c3781 100644 --- a/libs/auth/cache/file.go +++ b/libs/auth/cache/file.go @@ -73,7 +73,7 @@ func (c *FileTokenCache) Lookup(key string) (*oauth2.Token, error) { return t, nil } -func (c *FileTokenCache) DeleteKey(key string) error { +func (c *FileTokenCache) Delete(key string) error { err := c.load() if errors.Is(err, fs.ErrNotExist) { return ErrNotConfigured diff --git a/libs/auth/cache/file_test.go b/libs/auth/cache/file_test.go index fb7375c2..bd3f84ac 100644 --- a/libs/auth/cache/file_test.go +++ b/libs/auth/cache/file_test.go @@ -118,7 +118,7 @@ func TestStoreAndDeleteKey(t *testing.T) { require.NoError(t, err) l := &FileTokenCache{} - err = l.DeleteKey("x") + err = l.Delete("x") require.NoError(t, err) assert.Equal(t, 1, len(l.Tokens)) @@ -134,7 +134,7 @@ func TestDeleteKeyNotExist(t *testing.T) { c := &FileTokenCache{ Tokens: map[string]*oauth2.Token{}, } - err := c.DeleteKey("x") + err := c.Delete("x") assert.Equal(t, ErrNotConfigured, err) _, err = c.Lookup("x") diff --git a/libs/auth/cache/in_memory.go b/libs/auth/cache/in_memory.go index 6daaf868..756002e0 100644 --- a/libs/auth/cache/in_memory.go +++ b/libs/auth/cache/in_memory.go @@ -23,8 +23,8 @@ func (i *InMemoryTokenCache) Store(key string, t *oauth2.Token) error { return nil } -// DeleteKey implements TokenCache. -func (i *InMemoryTokenCache) DeleteKey(key string) error { +// Delete implements TokenCache. +func (i *InMemoryTokenCache) Delete(key string) error { _, ok := i.Tokens[key] if !ok { return ErrNotConfigured diff --git a/libs/auth/cache/in_memory_test.go b/libs/auth/cache/in_memory_test.go index 3714ed79..73995231 100644 --- a/libs/auth/cache/in_memory_test.go +++ b/libs/auth/cache/in_memory_test.go @@ -58,7 +58,7 @@ func TestInMemoryDeleteKey(t *testing.T) { }) require.NoError(t, err) - err = c.DeleteKey("x") + err = c.Delete("x") require.NoError(t, err) assert.Equal(t, 1, len(c.Tokens)) @@ -74,7 +74,7 @@ func TestInMemoryDeleteKeyNotExist(t *testing.T) { c := &InMemoryTokenCache{ Tokens: map[string]*oauth2.Token{}, } - err := c.DeleteKey("x") + err := c.Delete("x") assert.Equal(t, ErrNotConfigured, err) _, err = c.Lookup("x") diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go index 46f2af11..ec615fe7 100644 --- a/libs/auth/oauth.go +++ b/libs/auth/oauth.go @@ -152,7 +152,7 @@ func (a *PersistentAuth) ClearToken(ctx context.Context) error { } // lookup token identified by host (and possibly the account id) key := a.key() - return a.cache.DeleteKey(key) + return a.cache.Delete(key) } func (a *PersistentAuth) init(ctx context.Context) error { diff --git a/libs/auth/oauth_test.go b/libs/auth/oauth_test.go index e5e36032..fd245c67 100644 --- a/libs/auth/oauth_test.go +++ b/libs/auth/oauth_test.go @@ -72,7 +72,7 @@ func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) { return m.lookup(key) } -func (m *tokenCacheMock) DeleteKey(key string) error { +func (m *tokenCacheMock) Delete(key string) error { if m.deleteKey == nil { panic("no deleteKey mock") } From 89d3b1a4df8c3f5a4446169c5182f92e8ceb6801 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 20:23:09 +0200 Subject: [PATCH 08/13] remove redundant version specification --- libs/auth/cache/file.go | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/auth/cache/file.go b/libs/auth/cache/file.go index 4a7c3781..cca7d1a1 100644 --- a/libs/auth/cache/file.go +++ b/libs/auth/cache/file.go @@ -80,7 +80,6 @@ func (c *FileTokenCache) Delete(key string) error { } else if err != nil { return fmt.Errorf("load: %w", err) } - c.Version = tokenCacheVersion if c.Tokens == nil { c.Tokens = map[string]*oauth2.Token{} } From d037ec32a16c74e0fd5ee3da74fbb481d5deba49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 20:26:43 +0200 Subject: [PATCH 09/13] add new write function to persist to disk --- libs/auth/cache/file.go | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/libs/auth/cache/file.go b/libs/auth/cache/file.go index cca7d1a1..f1ac4c64 100644 --- a/libs/auth/cache/file.go +++ b/libs/auth/cache/file.go @@ -52,11 +52,7 @@ func (c *FileTokenCache) Store(key string, t *oauth2.Token) error { c.Tokens = map[string]*oauth2.Token{} } c.Tokens[key] = t - raw, err := json.MarshalIndent(c, "", " ") - if err != nil { - return fmt.Errorf("marshal: %w", err) - } - return os.WriteFile(c.fileLocation, raw, ownerReadWrite) + return c.write() } func (c *FileTokenCache) Lookup(key string) (*oauth2.Token, error) { @@ -88,11 +84,7 @@ func (c *FileTokenCache) Delete(key string) error { 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) + return c.write() } func (c *FileTokenCache) location() (string, error) { @@ -127,4 +119,12 @@ func (c *FileTokenCache) load() error { return nil } +func (c *FileTokenCache) write() error { + raw, err := json.MarshalIndent(c, "", " ") + if err != nil { + return fmt.Errorf("marshal: %w", err) + } + return os.WriteFile(c.fileLocation, raw, ownerReadWrite) +} + var _ TokenCache = (*FileTokenCache)(nil) From bb35ca090f80e1246fe4f342443ea34a72596e1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 20:37:53 +0200 Subject: [PATCH 10/13] logoutSession not exportable --- cmd/auth/logout.go | 40 ++++++++++++++++++++-------------------- cmd/auth/logout_test.go | 28 ++++++++++++++-------------- 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 37fba16f..3ed6594b 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -15,43 +15,43 @@ import ( "github.com/spf13/cobra" ) -type LogoutSession struct { - Profile string - File config.File - PersistentAuth *auth.PersistentAuth +type logoutSession struct { + profile string + file config.File + persistentAuth *auth.PersistentAuth } -func (l *LogoutSession) load(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth) error { - l.Profile = profileName - l.PersistentAuth = persistentAuth +func (l *logoutSession) load(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth) error { + l.profile = profileName + l.persistentAuth = persistentAuth iniFile, err := profile.DefaultProfiler.Get(ctx) if errors.Is(err, fs.ErrNotExist) { return err } else if err != nil { return fmt.Errorf("cannot parse config file: %w", err) } - l.File = *iniFile + l.file = *iniFile if err := l.setHostAndAccountIdFromProfile(); err != nil { return err } return nil } -func (l *LogoutSession) setHostAndAccountIdFromProfile() error { +func (l *logoutSession) setHostAndAccountIdFromProfile() error { sectionMap, err := l.getConfigSectionMap() if err != nil { return err } if sectionMap["host"] == "" { - return fmt.Errorf("no host configured for profile %s", l.Profile) + return fmt.Errorf("no host configured for profile %s", l.profile) } - l.PersistentAuth.Host = sectionMap["host"] - l.PersistentAuth.AccountID = sectionMap["account_id"] + l.persistentAuth.Host = sectionMap["host"] + l.persistentAuth.AccountID = sectionMap["account_id"] return nil } -func (l *LogoutSession) getConfigSectionMap() (map[string]string, error) { - section, err := l.File.GetSection(l.Profile) +func (l *logoutSession) getConfigSectionMap() (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) } @@ -59,16 +59,16 @@ func (l *LogoutSession) getConfigSectionMap() (map[string]string, error) { } // clear token from ~/.databricks/token-cache.json -func (l *LogoutSession) clearTokenCache(ctx context.Context) error { - return l.PersistentAuth.ClearToken(ctx) +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 *LogoutSession) 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, + ConfigFile: l.file.Path(), + Profile: l.profile, Host: sectionMap["host"], ClusterID: sectionMap["cluster_id"], WarehouseID: sectionMap["warehouse_id"], @@ -114,7 +114,7 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { } } defer persistentAuth.Close() - logoutSession := &LogoutSession{} + logoutSession := &logoutSession{} logoutSession.load(ctx, profileName, persistentAuth) configSectionMap, err := logoutSession.getConfigSectionMap() if err != nil { diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index 4f1133a3..d22002a5 100644 --- a/cmd/auth/logout_test.go +++ b/cmd/auth/logout_test.go @@ -26,11 +26,11 @@ func TestLogout_ClearConfigFile(t *testing.T) { require.NoError(t, err) iniFile, err := config.LoadFile(path) require.NoError(t, err) - logout := &LogoutSession{ - Profile: "abc", - File: *iniFile, + logout := &logoutSession{ + profile: "abc", + file: *iniFile, } - section, err := logout.File.GetSection("abc") + section, err := logout.file.GetSection("abc") assert.NoError(t, err) sectionMap := section.KeysHash() err = logout.clearConfigFile(ctx, sectionMap) @@ -62,15 +62,15 @@ func TestLogout_setHostAndAccountIdFromProfile(t *testing.T) { require.NoError(t, err) iniFile, err := config.LoadFile(path) require.NoError(t, err) - logout := &LogoutSession{ - Profile: "abc", - File: *iniFile, - PersistentAuth: &auth.PersistentAuth{}, + logout := &logoutSession{ + profile: "abc", + file: *iniFile, + persistentAuth: &auth.PersistentAuth{}, } err = logout.setHostAndAccountIdFromProfile() assert.NoError(t, err) - assert.Equal(t, logout.PersistentAuth.Host, "https://foo") - assert.Empty(t, logout.PersistentAuth.AccountID) + assert.Equal(t, logout.persistentAuth.Host, "https://foo") + assert.Empty(t, logout.persistentAuth.AccountID) } func TestLogout_getConfigSectionMap(t *testing.T) { @@ -86,10 +86,10 @@ func TestLogout_getConfigSectionMap(t *testing.T) { require.NoError(t, err) iniFile, err := config.LoadFile(path) require.NoError(t, err) - logout := &LogoutSession{ - Profile: "abc", - File: *iniFile, - PersistentAuth: &auth.PersistentAuth{}, + logout := &logoutSession{ + profile: "abc", + file: *iniFile, + persistentAuth: &auth.PersistentAuth{}, } configSectionMap, err := logout.getConfigSectionMap() assert.NoError(t, err) From b7ff019b6077ae8a2b65c28a0629f168b54af402 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 21:20:56 +0200 Subject: [PATCH 11/13] add test for file write --- libs/auth/cache/file_test.go | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/libs/auth/cache/file_test.go b/libs/auth/cache/file_test.go index bd3f84ac..48ab5c98 100644 --- a/libs/auth/cache/file_test.go +++ b/libs/auth/cache/file_test.go @@ -1,6 +1,7 @@ package cache import ( + "encoding/json" "os" "path/filepath" "runtime" @@ -140,3 +141,27 @@ func TestDeleteKeyNotExist(t *testing.T) { _, err = c.Lookup("x") assert.Equal(t, ErrNotConfigured, err) } + +func TestWrite(t *testing.T) { + tempFile := filepath.Join(t.TempDir(), "token-cache.json") + + tokenMap := map[string]*oauth2.Token{} + token := &oauth2.Token{ + AccessToken: "some-access-token", + } + tokenMap["test"] = token + + cache := &FileTokenCache{ + fileLocation: tempFile, + Tokens: tokenMap, + } + + err := cache.write() + assert.NoError(t, err) + + content, err := os.ReadFile(tempFile) + require.NoError(t, err) + + expected, _ := json.MarshalIndent(&cache, "", " ") + assert.Equal(t, content, expected) +} From 11c37673a6add53d5b29fa03fa7b1270985b2ff0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Mon, 23 Sep 2024 21:55:31 +0200 Subject: [PATCH 12/13] make tokenCacheMock consistent naming with struct --- libs/auth/oauth_test.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/libs/auth/oauth_test.go b/libs/auth/oauth_test.go index fd245c67..a8a1da70 100644 --- a/libs/auth/oauth_test.go +++ b/libs/auth/oauth_test.go @@ -53,9 +53,9 @@ func TestOidcForWorkspace(t *testing.T) { } type tokenCacheMock struct { - store func(key string, t *oauth2.Token) error - lookup func(key string) (*oauth2.Token, error) - deleteKey func(key string) error + store func(key string, t *oauth2.Token) error + lookup func(key string) (*oauth2.Token, error) + delete func(key string) error } func (m *tokenCacheMock) Store(key string, t *oauth2.Token) error { @@ -73,10 +73,10 @@ func (m *tokenCacheMock) Lookup(key string) (*oauth2.Token, error) { } func (m *tokenCacheMock) Delete(key string) error { - if m.deleteKey == nil { + if m.delete == nil { panic("no deleteKey mock") } - return m.deleteKey(key) + return m.delete(key) } func TestLoad(t *testing.T) { @@ -246,7 +246,7 @@ func TestClearToken(t *testing.T) { assert.Equal(t, "https://abc/oidc/accounts/xyz", key) return &oauth2.Token{}, ErrNotConfigured }, - deleteKey: func(key string) error { + delete: func(key string) error { assert.Equal(t, "https://abc/oidc/accounts/xyz", key) return nil }, @@ -269,7 +269,7 @@ func TestClearTokenNotExist(t *testing.T) { assert.Equal(t, "https://abc/oidc/accounts/xyz", key) return &oauth2.Token{}, ErrNotConfigured }, - deleteKey: func(key string) error { + delete: func(key string) error { assert.Equal(t, "https://abc/oidc/accounts/xyz", key) return ErrNotConfigured }, From 865964e0298b5eae46af3eee1e3e8d651a8c5575 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Nordstr=C3=B6m?= Date: Sun, 6 Oct 2024 23:37:05 +0200 Subject: [PATCH 13/13] reduce scope for logout cmd to only remove the OAuth token --- cmd/auth/logout.go | 37 ++++--------------------------------- cmd/auth/logout_test.go | 36 ------------------------------------ 2 files changed, 4 insertions(+), 69 deletions(-) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 3ed6594b..2153cfd7 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -9,7 +9,6 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/auth/cache" "github.com/databricks/cli/libs/cmdio" - "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" @@ -63,34 +62,11 @@ 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 *logoutSession) clearConfigFile(ctx context.Context, sectionMap map[string]string) error { - return databrickscfg.SaveToProfile(ctx, &config.Config{ - ConfigFile: l.file.Path(), - Profile: l.profile, - Host: sectionMap["host"], - ClusterID: sectionMap["cluster_id"], - WarehouseID: sectionMap["warehouse_id"], - ServerlessComputeID: sectionMap["serverless_compute_id"], - AccountID: sectionMap["account_id"], - Username: sectionMap["username"], - GoogleServiceAccount: sectionMap["google_service_account"], - AzureResourceID: sectionMap["azure_workspace_resource_id"], - AzureClientID: sectionMap["azure_client_id"], - AzureTenantID: sectionMap["azure_tenant_id"], - AzureEnvironment: sectionMap["azure_environment"], - AzureLoginAppID: sectionMap["azure_login_app_id"], - ClientID: sectionMap["client_id"], - AuthType: sectionMap["auth_type"], - }) -} - func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { cmd := &cobra.Command{ Use: "logout [PROFILE]", Short: "Logout from specified profile", - Long: "Clears OAuth token from token-cache and any sensitive value in the config file, if they exist.", + Long: "Removes the OAuth token from the token-cache", } cmd.RunE = func(cmd *cobra.Command, args []string) error { @@ -115,24 +91,19 @@ func newLogoutCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { } defer persistentAuth.Close() logoutSession := &logoutSession{} - logoutSession.load(ctx, profileName, persistentAuth) - configSectionMap, err := logoutSession.getConfigSectionMap() + err := logoutSession.load(ctx, profileName, persistentAuth) if err != nil { return err } err = logoutSession.clearTokenCache(ctx) if err != nil { if errors.Is(err, cache.ErrNotConfigured) { - // It is OK to not have OAuth configured. Move on and remove - // sensitive values from config file (Example PAT) + // It is OK to not have OAuth configured } else { return err } } - if err := logoutSession.clearConfigFile(ctx, configSectionMap); err != nil { - return err - } - cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully logged out", profileName)) + cmdio.LogString(ctx, fmt.Sprintf("Profile %s is logged out", profileName)) return nil } return cmd diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index d22002a5..27549274 100644 --- a/cmd/auth/logout_test.go +++ b/cmd/auth/logout_test.go @@ -13,42 +13,6 @@ import ( "github.com/databricks/databricks-sdk-go/config" ) -func TestLogout_ClearConfigFile(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", - Token: "xyz", - }) - require.NoError(t, err) - iniFile, err := config.LoadFile(path) - require.NoError(t, err) - logout := &logoutSession{ - profile: "abc", - file: *iniFile, - } - section, err := logout.file.GetSection("abc") - assert.NoError(t, err) - sectionMap := section.KeysHash() - err = logout.clearConfigFile(ctx, sectionMap) - assert.NoError(t, err) - - iniFile, err = config.LoadFile(path) - require.NoError(t, err) - assert.Len(t, iniFile.Sections(), 2) - assert.True(t, iniFile.HasSection("DEFAULT")) - assert.True(t, iniFile.HasSection("abc")) - - abc, err := iniFile.GetSection("abc") - assert.NoError(t, err) - raw := abc.KeysHash() - assert.Len(t, raw, 1) - assert.Equal(t, "https://foo", raw["host"]) -} - func TestLogout_setHostAndAccountIdFromProfile(t *testing.T) { ctx := context.Background() path := filepath.Join(t.TempDir(), "databrickscfg")