coder · Emyrk · Feb 26, 2024 · Feb 21, 2024 · Feb 22, 2024 · Feb 22, 2024
diff --git a/cli/organization.go b/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,13 +26,183 @@ func (r *RootCmd) organizations() *clibase.Cmd {
 		},
 		Children: []*clibase.Cmd{
 			r.currentOrganization(),
+			r.switchOrganization(),
 		},
 	}
 
 	cmd.Options = clibase.OptionSet{}
 	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)

diff --git a/cli/organization_test.go b/cli/organization_test.go
@@ -43,3 +43,37 @@ func TestCurrentOrganization(t *testing.T) {
 		pty.ExpectMatch(first.OrganizationID.String())
 	})
 }
+
+func TestOrganizationSwitch(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Switch", func(t *testing.T) {
+		t.Parallel()
+		ownerClient := coderdtest.New(t, nil)
+		first := coderdtest.CreateFirstUser(t, ownerClient)
+		// Owner is required to make orgs
+		client, _ := coderdtest.CreateAnotherUser(t, ownerClient, first.OrganizationID, rbac.RoleOwner())
+
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		orgs := []string{"foo", "bar"}
+		for _, orgName := range orgs {
+			_, err := client.CreateOrganization(ctx, codersdk.CreateOrganizationRequest{
+				Name: orgName,
+			})
+			require.NoError(t, err)
+		}
+
+		exp, err := client.OrganizationByName(ctx, "foo")
+		require.NoError(t, err)
+
+		inv, root := clitest.New(t, "organizations", "set", "foo")
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+		errC := make(chan error)
+		go func() {
+			errC <- inv.Run()
+		}()
+		require.NoError(t, <-errC)
+		pty.ExpectMatch(exp.ID.String())
+	})
+}
diff --git a/cli/root.go b/cli/root.go
@@ -733,7 +733,7 @@ func CurrentOrganization(r *RootCmd, inv *clibase.Invocation, client *codersdk.C
 		return org.IsDefault
 	})
 	if index < 0 {
-		return codersdk.Organization{}, xerrors.Errorf("unable to determine current organization. Use 'coder organizations switch <org>' to select an organization to use")
+		return codersdk.Organization{}, xerrors.Errorf("unable to determine current organization. Use 'coder set <org>' to select an organization to use")
 	}
 
 	return orgs[index], nil
@@ -1192,8 +1192,12 @@ func formatRunCommandError(err *clibase.RunCommandError, opts *formatOpts) strin
 func formatCoderSDKError(from string, err *codersdk.Error, opts *formatOpts) string {
 	var str strings.Builder
 	if opts.Verbose {
-		_, _ = str.WriteString(pretty.Sprint(headLineStyle(), fmt.Sprintf("API request error to \"%s:%s\". Status code %d", err.Method(), err.URL(), err.StatusCode())))
-		_, _ = str.WriteString("\n")
+		// If all these fields are empty, then do not print this information.
+		// This can occur if the error is being used outside the api.
+		if !(err.Method() == "" && err.URL() == "" && err.StatusCode() == 0) {
+			_, _ = str.WriteString(pretty.Sprint(headLineStyle(), fmt.Sprintf("API request error to \"%s:%s\". Status code %d", err.Method(), err.URL(), err.StatusCode())))
+			_, _ = str.WriteString("\n")
+		}
 	}
 	// Always include this trace. Users can ignore this.
 	if from != "" {