feat: add cli support for workspace automatic updates

coder · sreya · Nov 2, 2023 · Oct 31, 2023 · Nov 2, 2023 · Nov 2, 2023
commit cd7e9883e91cbff5dd8d2d301af97bbb1f1bdd1c
diff --git a/cli/autoupdate.go b/cli/autoupdate.go
@@ -0,0 +1,58 @@
+package cli
+
+import (
+	"fmt"
+	"strings"
+
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/cli/clibase"
+	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+func (r *RootCmd) autoupdate() *clibase.Cmd {
+	client := new(codersdk.Client)
+	cmd := &clibase.Cmd{
+		Annotations: workspaceCommand,
+		Use:         "autoupdate <workspace> <always|never>",
+		Short:       "Toggle auto-update policy for a workspace",
+		Middleware: clibase.Chain(
+			clibase.RequireNArgs(2),
+			r.InitClient(client),
+		),
+		Handler: func(inv *clibase.Invocation) error {
+			policy := strings.ToLower(inv.Args[1])
+			err := validateAutoUpdatePolicy(policy)
+			if err != nil {
+				return xerrors.Errorf("validate policy: %w", err)
+			}
+
+			workspace, err := namedWorkspace(inv.Context(), client, inv.Args[0])
+			if err != nil {
+				return xerrors.Errorf("get workspace: %w", err)
+			}
+
+			err = client.UpdateWorkspaceAutomaticUpdates(inv.Context(), workspace.ID, codersdk.UpdateWorkspaceAutomaticUpdatesRequest{
+				AutomaticUpdates: codersdk.AutomaticUpdates(policy),
+			})
+			if err != nil {
+				return xerrors.Errorf("update workspace automatic updates policy: %w", err)
+			}
+			_, _ = fmt.Fprintf(inv.Stdout, "Updated workspace %q auto-update policy to %q\n", workspace.Name, policy)
+			return nil
+		},
+	}
+
+	cmd.Options = append(cmd.Options, cliui.SkipPromptOption())
+	return cmd
+}
+
+func validateAutoUpdatePolicy(arg string) error {
+	switch codersdk.AutomaticUpdates(arg) {
+	case codersdk.AutomaticUpdatesAlways, codersdk.AutomaticUpdatesNever:
+		return nil
+	default:
+		return xerrors.Errorf("invalid option %q must be either of %q or %q", arg, codersdk.AutomaticUpdatesAlways, codersdk.AutomaticUpdatesNever)
+	}
+}
diff --git a/cli/autoupdate_test.go b/cli/autoupdate_test.go
@@ -0,0 +1,79 @@
+package cli_test
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+func TestAutoUpdate(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+		owner := coderdtest.CreateFirstUser(t, client)
+		member, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
+		version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, nil)
+		coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, member, owner.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+		require.Equal(t, codersdk.AutomaticUpdatesNever, workspace.AutomaticUpdates)
+
+		expectedPolicy := codersdk.AutomaticUpdatesAlways
+		inv, root := clitest.New(t, "autoupdate", workspace.Name, string(expectedPolicy))
+		clitest.SetupConfig(t, member, root)
+		var buf bytes.Buffer
+		inv.Stdout = &buf
+		err := inv.Run()
+		require.NoError(t, err)
+		require.Contains(t, buf.String(), fmt.Sprintf("Updated workspace %q auto-update policy to %q", workspace.Name, expectedPolicy))
+
+		workspace = coderdtest.MustWorkspace(t, client, workspace.ID)
+		require.Equal(t, expectedPolicy, workspace.AutomaticUpdates)
+	})
+
+	t.Run("InvalidArgs", func(t *testing.T) {
+		type testcase struct {
+			Name          string
+			Args          []string
+			ErrorContains string
+		}
+
+		cases := []testcase{
+			{
+				Name:          "NoPolicy",
+				Args:          []string{"autoupdate", "ws"},
+				ErrorContains: "wanted 2 args but got 1",
+			},
+			{
+				Name:          "InvalidPolicy",
+				Args:          []string{"autoupdate", "ws", "sometimes"},
+				ErrorContains: fmt.Sprintf("invalid option %q must be either of", "sometimes"),
+			},
+		}
+
+		for _, c := range cases {
+			c := c
+			t.Run(c.Name, func(t *testing.T) {
+				t.Parallel()
+				client := coderdtest.New(t, nil)
+				_ = coderdtest.CreateFirstUser(t, client)
+
+				inv, root := clitest.New(t, c.Args...)
+				clitest.SetupConfig(t, client, root)
+				err := inv.Run()
+				require.Error(t, err)
+				require.Contains(t, err.Error(), c.ErrorContains)
+			})
+		}
+	})
+}
diff --git a/cli/create.go b/cli/create.go
@@ -140,9 +140,9 @@ func (r *RootCmd) create() *clibase.Cmd {
 			}
 
 			richParameters, err := prepWorkspaceBuild(inv, client, prepWorkspaceBuildArgs{
-				Action:           WorkspaceCreate,
-				Template:         template,
-				NewWorkspaceName: workspaceName,
+				Action:            WorkspaceCreate,
+				TemplateVersionID: template.ActiveVersionID,
+				NewWorkspaceName:  workspaceName,
 
 				RichParameterFile: parameterFlags.richParameterFile,
 				RichParameters:    cliRichParameters,
@@ -224,10 +224,9 @@ func (r *RootCmd) create() *clibase.Cmd {
 }
 
 type prepWorkspaceBuildArgs struct {
-	Action           WorkspaceCLIAction
-	Template         codersdk.Template
-	NewWorkspaceName string
-	WorkspaceID      uuid.UUID
+	Action            WorkspaceCLIAction
+	TemplateVersionID uuid.UUID
+	NewWorkspaceName  string
 
 	LastBuildParameters []codersdk.WorkspaceBuildParameter
 
@@ -244,7 +243,7 @@ type prepWorkspaceBuildArgs struct {
 func prepWorkspaceBuild(inv *clibase.Invocation, client *codersdk.Client, args prepWorkspaceBuildArgs) ([]codersdk.WorkspaceBuildParameter, error) {
 	ctx := inv.Context()
 
-	templateVersion, err := client.TemplateVersion(ctx, args.Template.ActiveVersionID)
+	templateVersion, err := client.TemplateVersion(ctx, args.TemplateVersionID)
 	if err != nil {
 		return nil, xerrors.Errorf("get template version: %w", err)
 	}

diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go
@@ -600,9 +600,9 @@ func (r *RootCmd) scaletestCreateWorkspaces() *clibase.Cmd {
 			}
 
 			richParameters, err := prepWorkspaceBuild(inv, client, prepWorkspaceBuildArgs{
-				Action:           WorkspaceCreate,
-				Template:         tpl,
-				NewWorkspaceName: "scaletest-N", // TODO: the scaletest runner will pass in a different name here. Does this matter?
+				Action:            WorkspaceCreate,
+				TemplateVersionID: tpl.ActiveVersionID,
+				NewWorkspaceName:  "scaletest-N", // TODO: the scaletest runner will pass in a different name here. Does this matter?
 
 				RichParameterFile: parameterFlags.richParameterFile,
 				RichParameters:    cliRichParameters,

diff --git a/cli/parameter.go b/cli/parameter.go
@@ -20,6 +20,13 @@ type workspaceParameterFlags struct {
 
 	richParameterFile string
 	richParameters    []string
+
+	promptRichParameters bool
+}
+
+func (wpf *workspaceParameterFlags) allOptions() []clibase.Option {
+	options := append(wpf.cliBuildOptions(), wpf.cliParameters()...)
+	return append(options, wpf.alwaysPrompt())
 }
 
 func (wpf *workspaceParameterFlags) cliBuildOptions() []clibase.Option {
@@ -55,6 +62,14 @@ func (wpf *workspaceParameterFlags) cliParameters() []clibase.Option {
 	}
 }
 
+func (wpf *workspaceParameterFlags) alwaysPrompt() clibase.Option {
+	return clibase.Option{
+		Flag:        "always-prompt",
+		Description: "Always prompt all parameters. Does not pull parameter values from existing workspace.",
+		Value:       clibase.BoolOf(&wpf.promptRichParameters),
+	}
+}
+
 func asWorkspaceBuildParameters(nameValuePairs []string) ([]codersdk.WorkspaceBuildParameter, error) {
 	var params []codersdk.WorkspaceBuildParameter
 	for _, nameValue := range nameValuePairs {

diff --git a/cli/parameterresolver.go b/cli/parameterresolver.go
@@ -194,7 +194,7 @@ func (pr *ParameterResolver) resolveWithInput(resolved []codersdk.WorkspaceBuild
 			(action == WorkspaceUpdate && promptParameterOption) ||
 			(action == WorkspaceUpdate && tvp.Mutable && tvp.Required) ||
 			(action == WorkspaceUpdate && !tvp.Mutable && firstTimeUse) ||
-			(action == WorkspaceUpdate && tvp.Mutable && !tvp.Ephemeral && pr.promptRichParameters) {
+			(tvp.Mutable && !tvp.Ephemeral && pr.promptRichParameters) {
 			parameterValue, err := cliui.RichParameter(inv, tvp)
 			if err != nil {
 				return nil, err

diff --git a/cli/restart.go b/cli/restart.go
@@ -25,7 +25,7 @@ func (r *RootCmd) restart() *clibase.Cmd {
 			clibase.RequireNArgs(1),
 			r.InitClient(client),
 		),
-		Options: append(parameterFlags.cliBuildOptions(), cliui.SkipPromptOption()),
+		Options: clibase.OptionSet{cliui.SkipPromptOption()},
 		Handler: func(inv *clibase.Invocation) error {
 			ctx := inv.Context()
 			out := inv.Stdout
@@ -35,24 +35,10 @@ func (r *RootCmd) restart() *clibase.Cmd {
 				return err
 			}
 
-			lastBuildParameters, err := client.WorkspaceBuildParameters(inv.Context(), workspace.LatestBuild.ID)
-			if err != nil {
-				return err
-			}
-
-			buildOptions, err := asWorkspaceBuildParameters(parameterFlags.buildOptions)
-			if err != nil {
-				return xerrors.Errorf("can't parse build options: %w", err)
-			}
-
-			buildParameters, err := prepStartWorkspace(inv, client, prepStartWorkspaceArgs{
-				Action:            WorkspaceRestart,
-				TemplateVersionID: workspace.LatestBuild.TemplateVersionID,
-
-				LastBuildParameters: lastBuildParameters,
-
-				PromptBuildOptions: parameterFlags.promptBuildOptions,
-				BuildOptions:       buildOptions,
+			startReq, err := buildWorkspaceStartRequest(inv, client, startWorkspaceArgs{
+				workspace:      workspace,
+				action:         WorkspaceRestart,
+				parameterFlags: parameterFlags,
 			})
 			if err != nil {
 				return err
@@ -72,26 +58,21 @@ func (r *RootCmd) restart() *clibase.Cmd {
 			if err != nil {
 				return err
 			}
+
 			err = cliui.WorkspaceBuild(ctx, out, client, build.ID)
 			if err != nil {
 				return err
 			}
 
-			req := codersdk.CreateWorkspaceBuildRequest{
-				Transition:          codersdk.WorkspaceTransitionStart,
-				RichParameterValues: buildParameters,
-				TemplateVersionID:   workspace.LatestBuild.TemplateVersionID,
-			}
-
-			build, err = client.CreateWorkspaceBuild(ctx, workspace.ID, req)
+			build, err = client.CreateWorkspaceBuild(ctx, workspace.ID, startReq)
 			// It's possible for a workspace build to fail due to the template requiring starting
 			// workspaces with the active version.
 			if cerr, ok := codersdk.AsError(err); ok && cerr.StatusCode() == http.StatusUnauthorized {
-				build, err = startWorkspaceActiveVersion(inv, client, startWorkspaceActiveVersionArgs{
-					BuildOptions:        buildOptions,
-					LastBuildParameters: lastBuildParameters,
-					PromptBuildOptions:  parameterFlags.promptBuildOptions,
-					Workspace:           workspace,
+				_, _ = fmt.Fprintln(inv.Stdout, "Failed to restart with the template version from your last build. Policy may require you to restart with the current active template version.")
+				build, err = startWorkspace(inv, client, startWorkspaceArgs{
+					workspace:      workspace,
+					parameterFlags: parameterFlags,
+					action:         WorkspaceUpdate,
 				})
 				if err != nil {
 					return xerrors.Errorf("start workspace with active template version: %w", err)
@@ -112,5 +93,8 @@ func (r *RootCmd) restart() *clibase.Cmd {
 			return nil
 		},
 	}
+
+	cmd.Options = append(cmd.Options, parameterFlags.allOptions()...)
+
 	return cmd
 }
diff --git a/cli/restart_test.go b/cli/restart_test.go
@@ -239,4 +239,55 @@ func TestRestartWithParameters(t *testing.T) {
 			Value: immutableParameterValue,
 		})
 	})
+
+	t.Run("AlwaysPrompt", func(t *testing.T) {
+		t.Parallel()
+
+		// Create the workspace
+		client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+		owner := coderdtest.CreateFirstUser(t, client)
+		member, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
+		version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, mutableParamsResponse)
+		coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, member, owner.OrganizationID, template.ID, func(cwr *codersdk.CreateWorkspaceRequest) {
+			cwr.RichParameterValues = []codersdk.WorkspaceBuildParameter{
+				{
+					Name:  mutableParameterName,
+					Value: mutableParameterValue,
+				},
+			}
+		})
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+
+		inv, root := clitest.New(t, "restart", workspace.Name, "-y", "--always-prompt")
+		clitest.SetupConfig(t, member, root)
+		doneChan := make(chan struct{})
+		pty := ptytest.New(t).Attach(inv)
+		go func() {
+			defer close(doneChan)
+			err := inv.Run()
+			assert.NoError(t, err)
+		}()
+
+		// We should be prompted for the parameters again.
+		newValue := "xyz"
+		pty.ExpectMatch(mutableParameterName)
+		pty.WriteLine(newValue)
+		pty.ExpectMatch("workspace has been restarted")
+		<-doneChan
+
+		// Verify that the updated values are persisted.
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		defer cancel()
+
+		workspace, err := client.WorkspaceByOwnerAndName(ctx, workspace.OwnerName, workspace.Name, codersdk.WorkspaceOptions{})
+		require.NoError(t, err)
+		actualParameters, err := client.WorkspaceBuildParameters(ctx, workspace.LatestBuild.ID)
+		require.NoError(t, err)
+		require.Contains(t, actualParameters, codersdk.WorkspaceBuildParameter{
+			Name:  mutableParameterName,
+			Value: newValue,
+		})
+	})
 }
diff --git a/cli/root.go b/cli/root.go
@@ -97,6 +97,7 @@ func (r *RootCmd) Core() []*clibase.Cmd {
 		r.version(defaultVersionInfo),
 
 		// Workspace Commands
+		r.autoupdate(),
 		r.configSSH(),
 		r.create(),
 		r.deleteWorkspace(),