add logout cmd

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
 }
 

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

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