Add logout sub command

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(&perisistentAuth))
 	return cmd
 }
 

cmd/auth/logout.go

@@ -0,0 +1,139 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io/fs"
+
+	"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"
+)
+
+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
+	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
+	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
+}
+
+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)
+	}
+	return section.KeysHash(), nil
+}
+
+// clear token from ~/.databricks/token-cache.json
+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.",
+	}
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+		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
+			profileName, err = promptForProfile(ctx, persistentAuth.ProfileName())
+			if err != nil {
+				return err
+			}
+		}
+		defer persistentAuth.Close()
+		logoutSession := &LogoutSession{}
+		logoutSession.load(ctx, profileName, persistentAuth)
+		configSectionMap, err := logoutSession.getConfigSectionMap()
+		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)
+			} 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))
+		return nil
+	}
+	return cmd
+}

cmd/auth/logout_test.go

@@ -0,0 +1,98 @@
+package auth
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"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"
+)
+
+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")
+
+	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")
+}

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

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 {

libs/auth/cache/file_test.go

@@ -103,3 +103,40 @@ 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)
+}
+
+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.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)

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,40 @@ 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)
+}
+
+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)
+}