feat: set organization context in coder organizations (#12265)

* feat: add coder organizations set to change org context

`coder organizations set <org>`

cli/organization.go

@@ -1,12 +1,17 @@
 package cli
 
 import (
+	"errors"
 	"fmt"
+	"os"
+	"slices"
 	"strings"
 
 	"github.com/coder/coder/v2/cli/clibase"
 	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/cli/config"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/pretty"
 )
 
 func (r *RootCmd) organizations() *clibase.Cmd {
@@ -21,6 +26,7 @@ func (r *RootCmd) organizations() *clibase.Cmd {
 		},
 		Children: []*clibase.Cmd{
 			r.currentOrganization(),
+			r.switchOrganization(),
 		},
 	}
 
@@ -28,6 +34,175 @@ func (r *RootCmd) organizations() *clibase.Cmd {
 	return cmd
 }
 
+func (r *RootCmd) switchOrganization() *clibase.Cmd {
+	client := new(codersdk.Client)
+
+	cmd := &clibase.Cmd{
+		Use:   "set <organization name | ID>",
+		Short: "set the organization used by the CLI. Pass an empty string to reset to the default organization.",
+		Long: "set the organization used by the CLI. Pass an empty string to reset to the default organization.\n" + formatExamples(
+			example{
+				Description: "Remove the current organization and defer to the default.",
+				Command:     "coder organizations set ''",
+			},
+			example{
+				Description: "Switch to a custom organization.",
+				Command:     "coder organizations set my-org",
+			},
+		),
+		Middleware: clibase.Chain(
+			r.InitClient(client),
+			clibase.RequireRangeArgs(0, 1),
+		),
+		Options: clibase.OptionSet{},
+		Handler: func(inv *clibase.Invocation) error {
+			conf := r.createConfig()
+			orgs, err := client.OrganizationsByUser(inv.Context(), codersdk.Me)
+			if err != nil {
+				return fmt.Errorf("failed to get organizations: %w", err)
+			}
+			// Keep the list of orgs sorted
+			slices.SortFunc(orgs, func(a, b codersdk.Organization) int {
+				return strings.Compare(a.Name, b.Name)
+			})
+
+			var switchToOrg string
+			if len(inv.Args) == 0 {
+				// Pull switchToOrg from a prompt selector, rather than command line
+				// args.
+				switchToOrg, err = promptUserSelectOrg(inv, conf, orgs)
+				if err != nil {
+					return err
+				}
+			} else {
+				switchToOrg = inv.Args[0]
+			}
+
+			// If the user passes an empty string, we want to remove the organization
+			// from the config file. This will defer to default behavior.
+			if switchToOrg == "" {
+				err := conf.Organization().Delete()
+				if err != nil && !errors.Is(err, os.ErrNotExist) {
+					return fmt.Errorf("failed to unset organization: %w", err)
+				}
+				_, _ = fmt.Fprintf(inv.Stdout, "Organization unset\n")
+			} else {
+				// Find the selected org in our list.
+				index := slices.IndexFunc(orgs, func(org codersdk.Organization) bool {
+					return org.Name == switchToOrg || org.ID.String() == switchToOrg
+				})
+				if index < 0 {
+					// Using this error for better error message formatting
+					err := &codersdk.Error{
+						Response: codersdk.Response{
+							Message: fmt.Sprintf("Organization %q not found. Is the name correct, and are you a member of it?", switchToOrg),
+							Detail:  "Ensure the organization argument is correct and you are a member of it.",
+						},
+						Helper: fmt.Sprintf("Valid organizations you can switch to: %s", strings.Join(orgNames(orgs), ", ")),
+					}
+					return err
+				}
+
+				// Always write the uuid to the config file. Names can change.
+				err := conf.Organization().Write(orgs[index].ID.String())
+				if err != nil {
+					return fmt.Errorf("failed to write organization to config file: %w", err)
+				}
+			}
+
+			// Verify it worked.
+			current, err := CurrentOrganization(r, inv, client)
+			if err != nil {
+				// An SDK error could be a permission error. So offer the advice to unset the org
+				// and reset the context.
+				var sdkError *codersdk.Error
+				if errors.As(err, &sdkError) {
+					if sdkError.Helper == "" && sdkError.StatusCode() != 500 {
+						sdkError.Helper = `If this error persists, try unsetting your org with 'coder organizations set ""'`
+					}
+					return sdkError
+				}
+				return fmt.Errorf("failed to get current organization: %w", err)
+			}
+
+			_, _ = fmt.Fprintf(inv.Stdout, "Current organization context set to %s (%s)\n", current.Name, current.ID.String())
+			return nil
+		},
+	}
+
+	return cmd
+}
+
+// promptUserSelectOrg will prompt the user to select an organization from a list
+// of their organizations.
+func promptUserSelectOrg(inv *clibase.Invocation, conf config.Root, orgs []codersdk.Organization) (string, error) {
+	// Default choice
+	var defaultOrg string
+	// Comes from config file
+	if conf.Organization().Exists() {
+		defaultOrg, _ = conf.Organization().Read()
+	}
+
+	// No config? Comes from default org in the list
+	if defaultOrg == "" {
+		defIndex := slices.IndexFunc(orgs, func(org codersdk.Organization) bool {
+			return org.IsDefault
+		})
+		if defIndex >= 0 {
+			defaultOrg = orgs[defIndex].Name
+		}
+	}
+
+	// Defer to first org
+	if defaultOrg == "" && len(orgs) > 0 {
+		defaultOrg = orgs[0].Name
+	}
+
+	// Ensure the `defaultOrg` value is an org name, not a uuid.
+	// If it is a uuid, change it to the org name.
+	index := slices.IndexFunc(orgs, func(org codersdk.Organization) bool {
+		return org.ID.String() == defaultOrg || org.Name == defaultOrg
+	})
+	if index >= 0 {
+		defaultOrg = orgs[index].Name
+	}
+
+	// deselectOption is the option to delete the organization config file and defer
+	// to default behavior.
+	const deselectOption = "[Default]"
+	if defaultOrg == "" {
+		defaultOrg = deselectOption
+	}
+
+	// Pull value from a prompt
+	_, _ = fmt.Fprintln(inv.Stdout, pretty.Sprint(cliui.DefaultStyles.Wrap, "Select an organization below to set the current CLI context to:"))
+	value, err := cliui.Select(inv, cliui.SelectOptions{
+		Options:    append([]string{deselectOption}, orgNames(orgs)...),
+		Default:    defaultOrg,
+		Size:       10,
+		HideSearch: false,
+	})
+	if err != nil {
+		return "", err
+	}
+	// Deselect is an alias for ""
+	if value == deselectOption {
+		value = ""
+	}
+
+	return value, nil
+}
+
+// orgNames is a helper function to turn a list of organizations into a list of
+// their names as strings.
+func orgNames(orgs []codersdk.Organization) []string {
+	names := make([]string, 0, len(orgs))
+	for _, org := range orgs {
+		names = append(names, org.Name)
+	}
+	return names
+}
+
 func (r *RootCmd) currentOrganization() *clibase.Cmd {
 	var (
 		stringFormat func(orgs []codersdk.Organization) (string, error)