Add interactive profile picker to auth logout

cmd/auth/auth.go

@@ -30,6 +30,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`,
 
 	cmd.AddCommand(newEnvCommand())
 	cmd.AddCommand(newLoginCommand(&authArguments))
+	cmd.AddCommand(newLogoutCommand())
 	cmd.AddCommand(newProfilesCommand())
 	cmd.AddCommand(newTokenCommand(&authArguments))
 	cmd.AddCommand(newDescribeCommand())

cmd/auth/logout.go

@@ -0,0 +1,291 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"runtime"
+	"slices"
+	"strings"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/databrickscfg"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"github.com/manifoldco/promptui"
+	"github.com/spf13/cobra"
+)
+
+const logoutWarningTemplate = `{{ "Warning" | yellow }}: This will permanently log out of profile {{ .ProfileName | bold }}.
+
+The following changes will be made:
+  - Remove profile {{ .ProfileName | bold }} from {{ .ConfigPath }}
+  - Delete any cached OAuth tokens for this profile
+
+You will need to run {{ "databricks auth login" | bold }} to re-authenticate.
+`
+
+func newLogoutCommand() *cobra.Command {
+	defaultConfigPath := "~/.databrickscfg"
+	if runtime.GOOS == "windows" {
+		defaultConfigPath = "%USERPROFILE%\\.databrickscfg"
+	}
+
+	cmd := &cobra.Command{
+		Use:    "logout",
+		Short:  "Log out of a Databricks profile",
+		Hidden: true,
+		Long: fmt.Sprintf(`Log out of a Databricks profile.
+
+This command removes the specified profile from %s and deletes
+any associated cached OAuth tokens.
+
+You will need to run "databricks auth login" to re-authenticate after
+logging out.
+
+This command requires a profile to be specified (using --profile). If you
+omit --profile and run in an interactive terminal, you'll be shown an
+interactive profile picker to select which profile to log out of.
+
+While this command always removes the specified profile, the runtime behaviour
+depends on whether you run it in an interactive terminal and which flags you
+provide.
+
+1. If you specify --profile, the command will log out of that profile.
+   In an interactive terminal, you'll be asked to confirm unless --force
+   is specified.
+
+2. If you omit --profile and run in an interactive terminal, you'll be shown
+   an interactive picker listing all profiles from your configuration file.
+   Profiles are sorted alphabetically by name. You can search by profile
+   name, host, or account ID. After selecting a profile, you'll be asked to
+   confirm unless --force is specified.
+
+3. If you omit --profile and run in a non-interactive environment (e.g.
+   CI/CD pipelines), the command will fail with an error asking you to
+   specify --profile.
+
+4. Use --force to skip the confirmation prompt. This is required when
+   running in non-interactive mode; otherwise the command will fail.`,
+			defaultConfigPath),
+	}
+
+	var force bool
+	var profileName string
+	cmd.Flags().BoolVar(&force, "force", false, "Skip confirmation prompt")
+	cmd.Flags().StringVar(&profileName, "profile", "", "The profile to log out of")
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+
+		if profileName == "" {
+			if !cmdio.IsPromptSupported(ctx) {
+				return errors.New("the command is being run in a non-interactive environment, please specify a profile to log out of using --profile")
+			}
+			selected, err := promptForLogoutProfile(ctx, profile.DefaultProfiler)
+			if err != nil {
+				return err
+			}
+			profileName = selected
+		}
+
+		tokenCache, err := cache.NewFileTokenCache()
+		if err != nil {
+			log.Warnf(ctx, "Failed to open token cache: %v", err)
+		}
+
+		return runLogout(ctx, logoutArgs{
+			profileName:    profileName,
+			force:          force,
+			profiler:       profile.DefaultProfiler,
+			tokenCache:     tokenCache,
+			configFilePath: os.Getenv("DATABRICKS_CONFIG_FILE"),
+		})
+	}
+
+	return cmd
+}
+
+type logoutArgs struct {
+	profileName    string
+	force          bool
+	profiler       profile.Profiler
+	tokenCache     cache.TokenCache
+	configFilePath string
+}
+
+func runLogout(ctx context.Context, args logoutArgs) error {
+	matchedProfile, err := getMatchingProfile(ctx, args.profileName, args.profiler)
+	if err != nil {
+		return err
+	}
+
+	if !args.force {
+		if !cmdio.IsPromptSupported(ctx) {
+			return errors.New("please specify --force to skip confirmation in non-interactive mode")
+		}
+
+		configPath, err := args.profiler.GetPath(ctx)
+		if err != nil {
+			return err
+		}
+		err = cmdio.RenderWithTemplate(ctx, map[string]string{
+			"ProfileName": args.profileName,
+			"ConfigPath":  configPath,
+		}, "", logoutWarningTemplate)
+		if err != nil {
+			return err
+		}
+
+		approved, err := cmdio.AskYesOrNo(ctx, "Are you sure?")
+		if err != nil {
+			return err
+		}
+		if !approved {
+			return nil
+		}
+	}
+
+	// First delete the profile and then perform best-effort token cache cleanup
+	// to avoid partial cleanup in case of errors from profile deletion.
+	err = databrickscfg.DeleteProfile(ctx, args.profileName, args.configFilePath)
+	if err != nil {
+		return fmt.Errorf("failed to remove profile: %w", err)
+	}
+
+	clearTokenCache(ctx, *matchedProfile, args.profiler, args.tokenCache)
+
+	cmdio.LogString(ctx, fmt.Sprintf("Successfully logged out of profile %q.", args.profileName))
+	return nil
+}
+
+// getMatchingProfile loads a profile by name and returns an error with
+// available profile names if the profile is not found.
+func getMatchingProfile(ctx context.Context, profileName string, profiler profile.Profiler) (*profile.Profile, error) {
+	profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName))
+	if err != nil {
+		return nil, err
+	}
+
+	if len(profiles) == 0 {
+		allProfiles, err := profiler.LoadProfiles(ctx, profile.MatchAllProfiles)
+		if err != nil {
+			return nil, fmt.Errorf("profile %q not found", profileName)
+		}
+
+		return nil, fmt.Errorf("profile %q not found. Available profiles: %s", profileName, allProfiles.Names())
+	}
+
+	return &profiles[0], nil
+}
+
+type logoutProfileItem struct {
+	PaddedName string
+	profile.Profile
+}
+
+// promptForLogoutProfile shows an interactive profile picker for logout.
+// Account profiles are displayed as "name (account: id)", workspace profiles
+// as "name (host)".
+func promptForLogoutProfile(ctx context.Context, profiler profile.Profiler) (string, error) {
+	allProfiles, err := profiler.LoadProfiles(ctx, profile.MatchAllProfiles)
+	if err != nil {
+		return "", err
+	}
+	if len(allProfiles) == 0 {
+		return "", errors.New("no profiles configured. Run 'databricks auth login' to create a profile")
+	}
+
+	slices.SortFunc(allProfiles, func(a, b profile.Profile) int {
+		return strings.Compare(strings.ToLower(a.Name), strings.ToLower(b.Name))
+	})
+
+	maxNameLen := 0
+	for _, p := range allProfiles {
+		maxNameLen = max(maxNameLen, len(p.Name))
+	}
+
+	items := make([]logoutProfileItem, len(allProfiles))
+	for i, p := range allProfiles {
+		items[i] = logoutProfileItem{
+			PaddedName: fmt.Sprintf("%-*s", maxNameLen, p.Name),
+			Profile:    p,
+		}
+	}
+
+	i, _, err := cmdio.RunSelect(ctx, &promptui.Select{
+		Label:             "Select a profile to log out of",
+		Items:             items,
+		StartInSearchMode: len(items) > 5,
+		// Allow searching by name, host, and account ID.
+		Searcher: func(input string, index int) bool {
+			input = strings.ToLower(input)
+			name := strings.ToLower(items[index].Name)
+			host := strings.ToLower(items[index].Host)
+			accountID := strings.ToLower(items[index].AccountID)
+			return strings.Contains(name, input) || strings.Contains(host, input) || strings.Contains(accountID, input)
+		},
+		Templates: &promptui.SelectTemplates{
+			Label:    "{{ . | faint }}",
+			Active:   `▸ {{.PaddedName | bold}}{{if .AccountID}} (account: {{.AccountID}}){{else}} ({{.Host}}){{end}}`,
+			Inactive: `  {{.PaddedName}}{{if .AccountID}} (account: {{.AccountID | faint}}){{else}} ({{.Host | faint}}){{end}}`,
+			Selected: `{{ "Selected profile" | faint }}: {{ .Name | bold }}`,
+		},
+	})
+	if err != nil {
+		return "", err
+	}
+	return items[i].Name, nil
+}
+
+// clearTokenCache removes cached OAuth tokens for the given profile from the
+// token cache. It removes:
+//  1. The entry keyed by the profile name.
+//  2. The entry keyed by the host-based cache key, but only if no other
+//     remaining profile references the same key. For account and unified
+//     profiles, the cache key includes the OIDC path
+//     (host/oidc/accounts/<account_id>).
+func clearTokenCache(ctx context.Context, p profile.Profile, profiler profile.Profiler, tokenCache cache.TokenCache) {
+	if tokenCache == nil {
+		return
+	}
+
+	if err := tokenCache.Store(p.Name, nil); err != nil {
+		log.Warnf(ctx, "Failed to delete profile-keyed token for profile %q: %v", p.Name, err)
+	}
+
+	hostCacheKey, matchFn := hostCacheKeyAndMatchFn(p)
+	if hostCacheKey == "" {
+		return
+	}
+
+	otherProfiles, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool {
+		return candidate.Name != p.Name && matchFn(candidate)
+	})
+	if err != nil {
+		log.Warnf(ctx, "Failed to load profiles for host cache key %q: %v", hostCacheKey, err)
+		return
+	}
+
+	if len(otherProfiles) == 0 {
+		if err := tokenCache.Store(hostCacheKey, nil); err != nil {
+			log.Warnf(ctx, "Failed to delete host-keyed token for %q: %v", hostCacheKey, err)
+		}
+	}
+}
+
+// hostCacheKeyAndMatchFn returns the token cache key and a profile match
+// function for the host-based token entry. Account and unified profiles use
+// host/oidc/accounts/<account_id> as the cache key and match on both host and
+// account ID; workspace profiles use just the host.
+func hostCacheKeyAndMatchFn(p profile.Profile) (string, profile.ProfileMatchFunction) {
+	host := strings.TrimRight(p.Host, "/")
+
+	if p.AccountID != "" {
+		return host + "/oidc/accounts/" + p.AccountID, profile.WithHostAndAccountID(host, p.AccountID)
+	}
+
+	return host, profile.WithHost(host)
+}