coder · ethanndickson · Aug 20, 2024 · Aug 9, 2024
diff --git a/cli/cliui/output.go b/cli/cliui/output.go
@@ -65,8 +65,8 @@ func (f *OutputFormatter) AttachOptions(opts *serpent.OptionSet) {
 			Flag:          "output",
 			FlagShorthand: "o",
 			Default:       f.formats[0].ID(),
-			Value:         serpent.StringOf(&f.formatID),
-			Description:   "Output format. Available formats: " + strings.Join(formatNames, ", ") + ".",
+			Value:         serpent.EnumOf(&f.formatID, formatNames...),
+			Description:   "Output format.",
 		},
 	)
 }
@@ -136,8 +136,8 @@ func (f *tableFormat) AttachOptions(opts *serpent.OptionSet) {
 			Flag:          "column",
 			FlagShorthand: "c",
 			Default:       strings.Join(f.defaultColumns, ","),
-			Value:         serpent.StringArrayOf(&f.columns),
-			Description:   "Columns to display in table output. Available columns: " + strings.Join(f.allColumns, ", ") + ".",
+			Value:         serpent.EnumArrayOf(&f.columns, f.allColumns...),
+			Description:   "Columns to display in table output.",
 		},
 	)
 }

diff --git a/cli/cliui/output_test.go b/cli/cliui/output_test.go
@@ -106,11 +106,11 @@ func Test_OutputFormatter(t *testing.T) {
 
 		fs := cmd.Options.FlagSet()
 
-		selected, err := fs.GetString("output")
-		require.NoError(t, err)
-		require.Equal(t, "json", selected)
+		selected := cmd.Options.ByFlag("output")
+		require.NotNil(t, selected)
+		require.Equal(t, "json", selected.Value.String())
 		usage := fs.FlagUsages()
-		require.Contains(t, usage, "Available formats: json, foo")
+		require.Contains(t, usage, "Output format.")
 		require.Contains(t, usage, "foo flag 1234")
 
 		ctx := context.Background()
@@ -129,11 +129,10 @@ func Test_OutputFormatter(t *testing.T) {
 		require.Equal(t, "foo", out)
 		require.EqualValues(t, 1, atomic.LoadInt64(&called))
 
-		require.NoError(t, fs.Set("output", "bar"))
+		require.Error(t, fs.Set("output", "bar"))
 		out, err = f.Format(ctx, data)
-		require.Error(t, err)
-		require.ErrorContains(t, err, "bar")
-		require.Equal(t, "", out)
-		require.EqualValues(t, 1, atomic.LoadInt64(&called))
+		require.NoError(t, err)
+		require.Equal(t, "foo", out)
+		require.EqualValues(t, 2, atomic.LoadInt64(&called))
 	})
 }
diff --git a/cli/completion.go b/cli/completion.go
@@ -0,0 +1,89 @@
+package cli
+
+import (
+	"fmt"
+
+	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/serpent"
+	"github.com/coder/serpent/completion"
+)
+
+func (*RootCmd) completion() *serpent.Command {
+	var shellName string
+	var printOutput bool
+	shellOptions := completion.ShellOptions(&shellName)
+	return &serpent.Command{
+		Use:   "completion",
+		Short: "Install or update shell completion scripts for the detected or chosen shell.",
+		Options: []serpent.Option{
+			{
+				Flag:          "shell",
+				FlagShorthand: "s",
+				Description:   "The shell to install completion for.",
+				Value:         shellOptions,
+			},
+			{
+				Flag:          "print",
+				Description:   "Print the completion script instead of installing it.",
+				FlagShorthand: "p",
+
+				Value: serpent.BoolOf(&printOutput),
+			},
+		},
+		Handler: func(inv *serpent.Invocation) error {
+			if shellName != "" {
+				shell, err := completion.ShellByName(shellName, inv.Command.Parent.Name())
+				if err != nil {
+					return err
+				}
+				if printOutput {
+					return shell.WriteCompletion(inv.Stdout)
+				}
+				return installCompletion(inv, shell)
+			}
+			shell, err := completion.DetectUserShell(inv.Command.Parent.Name())
+			if err == nil {
+				return installCompletion(inv, shell)
+			}
+			// Silently continue to the shell selection if detecting failed.
+			choice, err := cliui.Select(inv, cliui.SelectOptions{
+				Message: "Select a shell to install completion for:",
+				Options: shellOptions.Choices,
+			})
+			if err != nil {
+				return err
+			}
+			shellChoice, err := completion.ShellByName(choice, inv.Command.Parent.Name())
+			if err != nil {
+				return err
+			}
+			if printOutput {
+				return shellChoice.WriteCompletion(inv.Stdout)
+			}
+			return installCompletion(inv, shellChoice)
+		},
+	}
+}
+
+func installCompletion(inv *serpent.Invocation, shell completion.Shell) error {
+	path, err := shell.InstallPath()
+	if err != nil {
+		cliui.Error(inv.Stderr, fmt.Sprintf("Failed to determine completion path %v", err))
+		return shell.WriteCompletion(inv.Stdout)
+	}
+	choice, err := cliui.Select(inv, cliui.SelectOptions{
+		Options: []string{
+			"Confirm",
+			"Print to terminal",
+		},
+		Message:    fmt.Sprintf("Install completion for %s at %s?", shell.Name(), path),
+		HideSearch: true,
+	})
+	if err != nil {
+		return err
+	}
+	if choice == "Print to terminal" {
+		return shell.WriteCompletion(inv.Stdout)
+	}
+	return completion.InstallShellCompletion(shell)
+}
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -17,6 +17,7 @@ import (
 	"strings"
 
 	"github.com/cli/safeexec"
+	"github.com/natefinch/atomic"
 	"github.com/pkg/diff"
 	"github.com/pkg/diff/write"
 	"golang.org/x/exp/constraints"
@@ -524,7 +525,7 @@ func (r *RootCmd) configSSH() *serpent.Command {
 			}
 
 			if !bytes.Equal(configRaw, configModified) {
-				err = writeWithTempFileAndMove(sshConfigFile, bytes.NewReader(configModified))
+				err = atomic.WriteFile(sshConfigFile, bytes.NewReader(configModified))
 				if err != nil {
 					return xerrors.Errorf("write ssh config failed: %w", err)
 				}
@@ -758,50 +759,6 @@ func sshConfigSplitOnCoderSection(data []byte) (before, section []byte, after []
 	return data, nil, nil, nil
 }
 
-// writeWithTempFileAndMove writes to a temporary file in the same
-// directory as path and renames the temp file to the file provided in
-// path. This ensure we avoid trashing the file we are writing due to
-// unforeseen circumstance like filesystem full, command killed, etc.
-func writeWithTempFileAndMove(path string, r io.Reader) (err error) {
-	dir := filepath.Dir(path)
-	name := filepath.Base(path)
-
-	// Ensure that e.g. the ~/.ssh directory exists.
-	if err = os.MkdirAll(dir, 0o700); err != nil {
-		return xerrors.Errorf("create directory: %w", err)
-	}
-
-	// Create a tempfile in the same directory for ensuring write
-	// operation does not fail.
-	f, err := os.CreateTemp(dir, fmt.Sprintf(".%s.", name))
-	if err != nil {
-		return xerrors.Errorf("create temp file failed: %w", err)
-	}
-	defer func() {
-		if err != nil {
-			_ = os.Remove(f.Name()) // Cleanup in case a step failed.
-		}
-	}()
-
-	_, err = io.Copy(f, r)
-	if err != nil {
-		_ = f.Close()
-		return xerrors.Errorf("write temp file failed: %w", err)
-	}
-
-	err = f.Close()
-	if err != nil {
-		return xerrors.Errorf("close temp file failed: %w", err)
-	}
-
-	err = os.Rename(f.Name(), path)
-	if err != nil {
-		return xerrors.Errorf("rename temp file failed: %w", err)
-	}
-
-	return nil
-}
-
 // sshConfigExecEscape quotes the string if it contains spaces, as per
 // `man 5 ssh_config`. However, OpenSSH uses exec in the users shell to
 // run the command, and as such the formatting/escape requirements

diff --git a/cli/help.go b/cli/help.go
@@ -81,6 +81,8 @@ var usageTemplate = func() *template.Template {
 					switch v := opt.Value.(type) {
 					case *serpent.Enum:
 						return strings.Join(v.Choices, "|")
+					case *serpent.EnumArray:
+						return fmt.Sprintf("[%s]", strings.Join(v.Choices, "|"))
 					default:
 						return v.Type()
 					}

diff --git a/cli/organizationmembers.go b/cli/organizationmembers.go
@@ -137,7 +137,7 @@ func (r *RootCmd) assignOrganizationRoles(orgContext *OrganizationContext) *serp
 
 func (r *RootCmd) listOrganizationMembers(orgContext *OrganizationContext) *serpent.Command {
 	formatter := cliui.NewOutputFormatter(
-		cliui.TableFormat([]codersdk.OrganizationMemberWithUserData{}, []string{"username", "organization_roles"}),
+		cliui.TableFormat([]codersdk.OrganizationMemberWithUserData{}, []string{"username", "organization roles"}),
 		cliui.JSONFormat(),
 	)
 

diff --git a/cli/organizationmembers_test.go b/cli/organizationmembers_test.go
@@ -23,7 +23,7 @@ func TestListOrganizationMembers(t *testing.T) {
 		client, user := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin())
 
 		ctx := testutil.Context(t, testutil.WaitMedium)
-		inv, root := clitest.New(t, "organization", "members", "list", "-c", "user_id,username,roles")
+		inv, root := clitest.New(t, "organization", "members", "list", "-c", "user id,username,organization roles")
 		clitest.SetupConfig(t, client, root)
 
 		buf := new(bytes.Buffer)

diff --git a/cli/organizationroles.go b/cli/organizationroles.go
@@ -36,7 +36,7 @@ func (r *RootCmd) organizationRoles(orgContext *OrganizationContext) *serpent.Co
 func (r *RootCmd) showOrganizationRoles(orgContext *OrganizationContext) *serpent.Command {
 	formatter := cliui.NewOutputFormatter(
 		cliui.ChangeFormatterData(
-			cliui.TableFormat([]roleTableRow{}, []string{"name", "display_name", "site_permissions", "organization_permissions", "user_permissions"}),
+			cliui.TableFormat([]roleTableRow{}, []string{"name", "display name", "site permissions", "organization permissions", "user permissions"}),
 			func(data any) (any, error) {
 				inputs, ok := data.([]codersdk.AssignableRoles)
 				if !ok {
@@ -103,7 +103,7 @@ func (r *RootCmd) showOrganizationRoles(orgContext *OrganizationContext) *serpen
 func (r *RootCmd) editOrganizationRole(orgContext *OrganizationContext) *serpent.Command {
 	formatter := cliui.NewOutputFormatter(
 		cliui.ChangeFormatterData(
-			cliui.TableFormat([]roleTableRow{}, []string{"name", "display_name", "site_permissions", "organization_permissions", "user_permissions"}),
+			cliui.TableFormat([]roleTableRow{}, []string{"name", "display name", "site permissions", "organization permissions", "user permissions"}),
 			func(data any) (any, error) {
 				typed, _ := data.(codersdk.Role)
 				return []roleTableRow{roleToTableView(typed)}, nil
@@ -408,10 +408,10 @@ func roleToTableView(role codersdk.Role) roleTableRow {
 
 type roleTableRow struct {
 	Name            string `table:"name,default_sort"`
-	DisplayName     string `table:"display_name"`
-	OrganizationID  string `table:"organization_id"`
-	SitePermissions string ` table:"site_permissions"`
+	DisplayName     string `table:"display name"`
+	OrganizationID  string `table:"organization id"`
+	SitePermissions string ` table:"site permissions"`
 	// map[<org_id>] -> Permissions
-	OrganizationPermissions string `table:"organization_permissions"`
-	UserPermissions         string `table:"user_permissions"`
+	OrganizationPermissions string `table:"organization permissions"`
+	UserPermissions         string `table:"user permissions"`
 }
diff --git a/cli/root.go b/cli/root.go
@@ -82,6 +82,7 @@ const (
 func (r *RootCmd) CoreSubcommands() []*serpent.Command {
 	// Please re-sort this list alphabetically if you change it!
 	return []*serpent.Command{
+		r.completion(),
 		r.dotfiles(),
 		r.externalAuth(),
 		r.login(),

diff --git a/cli/stat.go b/cli/stat.go
@@ -32,11 +32,11 @@ func (r *RootCmd) stat() *serpent.Command {
 		fs        = afero.NewReadOnlyFs(afero.NewOsFs())
 		formatter = cliui.NewOutputFormatter(
 			cliui.TableFormat([]statsRow{}, []string{
-				"host_cpu",
-				"host_memory",
-				"home_disk",
-				"container_cpu",
-				"container_memory",
+				"host cpu",
+				"host memory",
+				"home disk",
+				"container cpu",
+				"container memory",
 			}),
 			cliui.JSONFormat(),
 		)
@@ -284,9 +284,9 @@ func (*RootCmd) statDisk(fs afero.Fs) *serpent.Command {
 }
 
 type statsRow struct {
-	HostCPU         *clistat.Result `json:"host_cpu" table:"host_cpu,default_sort"`
-	HostMemory      *clistat.Result `json:"host_memory" table:"host_memory"`
-	Disk            *clistat.Result `json:"home_disk" table:"home_disk"`
-	ContainerCPU    *clistat.Result `json:"container_cpu" table:"container_cpu"`
-	ContainerMemory *clistat.Result `json:"container_memory" table:"container_memory"`
+	HostCPU         *clistat.Result `json:"host_cpu" table:"host cpu,default_sort"`
+	HostMemory      *clistat.Result `json:"host_memory" table:"host memory"`
+	Disk            *clistat.Result `json:"home_disk" table:"home disk"`
+	ContainerCPU    *clistat.Result `json:"container_cpu" table:"container cpu"`
+	ContainerMemory *clistat.Result `json:"container_memory" table:"container memory"`
 }
diff --git a/cli/templateedit.go b/cli/templateedit.go
@@ -3,7 +3,6 @@ package cli
 import (
 	"fmt"
 	"net/http"
-	"strings"
 	"time"
 
 	"golang.org/x/xerrors"
@@ -239,35 +238,14 @@ func (r *RootCmd) templateEdit() *serpent.Command {
 			Value:       serpent.DurationOf(&activityBump),
 		},
 		{
-			Flag: "autostart-requirement-weekdays",
-			// workspaces created from this template must be restarted on the given weekdays. To unset this value for the template (and disable the autostop requirement for the template), pass 'none'.
+			Flag:        "autostart-requirement-weekdays",
 			Description: "Edit the template autostart requirement weekdays - workspaces created from this template can only autostart on the given weekdays. To unset this value for the template (and allow autostart on all days), pass 'all'.",
-			Value: serpent.Validate(serpent.StringArrayOf(&autostartRequirementDaysOfWeek), func(value *serpent.StringArray) error {
-				v := value.GetSlice()
-				if len(v) == 1 && v[0] == "all" {
-					return nil
-				}
-				_, err := codersdk.WeekdaysToBitmap(v)
-				if err != nil {
-					return xerrors.Errorf("invalid autostart requirement days of week %q: %w", strings.Join(v, ","), err)
-				}
-				return nil
-			}),
+			Value:       serpent.EnumArrayOf(&autostartRequirementDaysOfWeek, append(codersdk.AllDaysOfWeek, "all")...),
 		},
 		{
 			Flag:        "autostop-requirement-weekdays",
 			Description: "Edit the template autostop requirement weekdays - workspaces created from this template must be restarted on the given weekdays. To unset this value for the template (and disable the autostop requirement for the template), pass 'none'.",
-			Value: serpent.Validate(serpent.StringArrayOf(&autostopRequirementDaysOfWeek), func(value *serpent.StringArray) error {
-				v := value.GetSlice()
-				if len(v) == 1 && v[0] == "none" {
-					return nil
-				}
-				_, err := codersdk.WeekdaysToBitmap(v)
-				if err != nil {
-					return xerrors.Errorf("invalid autostop requirement days of week %q: %w", strings.Join(v, ","), err)
-				}
-				return nil
-			}),
+			Value:       serpent.EnumArrayOf(&autostopRequirementDaysOfWeek, append(codersdk.AllDaysOfWeek, "none")...),
 		},
 		{
 			Flag:        "autostop-requirement-weeks",

diff --git a/cli/templateversions.go b/cli/templateversions.go
@@ -40,11 +40,11 @@ func (r *RootCmd) templateVersions() *serpent.Command {
 
 func (r *RootCmd) templateVersionsList() *serpent.Command {
 	defaultColumns := []string{
-		"Name",
-		"Created At",
-		"Created By",
-		"Status",
-		"Active",
+		"name",
+		"created at",
+		"created by",
+		"status",
+		"active",
 	}
 	formatter := cliui.NewOutputFormatter(
 		cliui.TableFormat([]templateVersionRow{}, defaultColumns),
@@ -70,10 +70,10 @@ func (r *RootCmd) templateVersionsList() *serpent.Command {
 						for _, opt := range i.Command.Options {
 							if opt.Flag == "column" {
 								if opt.ValueSource == serpent.ValueSourceDefault {
-									v, ok := opt.Value.(*serpent.StringArray)
+									v, ok := opt.Value.(*serpent.EnumArray)
 									if ok {
 										// Add the extra new default column.
-										*v = append(*v, "Archived")
+										_ = v.Append("Archived")
 									}
 								}
 								break

diff --git a/cli/testdata/coder_--help.golden b/cli/testdata/coder_--help.golden
@@ -15,6 +15,8 @@ USAGE:
 
 SUBCOMMANDS:
     autoupdate        Toggle auto-update policy for a workspace
+    completion        Install or update shell completion scripts for the
+                      detected or chosen shell.
     config-ssh        Add an SSH Host entry for your workspaces "ssh
                       coder.workspace"
     create            Create a workspace