refactor: simplify select current workspace task (#117)

internal/module/service.go

@@ -251,17 +251,13 @@ func (s *Service) GetByPath(path string) (*Module, error) {
 	return nil, fmt.Errorf("%s: %w", path, resource.ErrNotFound)
 }
 
+// SetCurrent sets the current workspace for the module.
 func (s *Service) SetCurrent(moduleID, workspaceID resource.ID) error {
 	_, err := s.table.Update(moduleID, func(existing *Module) error {
 		existing.CurrentWorkspaceID = &workspaceID
 		return nil
 	})
-	if err != nil {
-		s.logger.Error("setting current workspace", "module", moduleID, "workspace", workspaceID, "error", err)
-		return err
-	}
-	s.logger.Debug("set current workspace", "module", moduleID, "workspace", workspaceID)
-	return nil
+	return err
 }
 
 // updateSpec updates the task spec with common module settings.

internal/workspace/service.go

@@ -259,15 +259,6 @@ func (s *Service) List(opts ListOptions) []*Workspace {
 // command, which sets the current workspace for the module. Once that's
 // finished it then updates the current workspace in pug itself too.
 func (s *Service) SelectWorkspace(moduleID, workspaceID resource.ID) error {
-	if err := s.selectWorkspace(moduleID, workspaceID); err != nil {
-		s.logger.Error("selecting current workspace", "workspace_id", workspaceID, "error", err)
-		return err
-	}
-	s.logger.Debug("selected current workspace", "workspace", workspaceID)
-	return nil
-}
-
-func (s *Service) selectWorkspace(moduleID, workspaceID resource.ID) error {
 	ws, err := s.table.Get(workspaceID)
 	if err != nil {
 		return err
@@ -284,15 +275,16 @@ func (s *Service) selectWorkspace(moduleID, workspaceID resource.ID) error {
 		Args:      []string{ws.Name},
 		Immediate: true,
 		Wait:      true,
+		BeforeExited: func(t *task.Task) (task.Summary, error) {
+			// Now the terraform command has finished, update the current
+			// workspace in pug as well.
+			err := s.modules.SetCurrent(moduleID, workspaceID)
+			return nil, err
+		},
 	})
 	if err != nil {
 		return err
 	}
-	// Now task has finished successfully, update the current workspace in pug
-	// as well.
-	if err := s.modules.SetCurrent(moduleID, workspaceID); err != nil {
-		return err
-	}
 	return nil
 }