Merge pull request #140 from trickest/feat/stop-cmd-status-flag

Add status flag to the stop command

Author: PolovinaD

README.md

@@ -162,7 +162,7 @@ More example workflow **config.yaml** files can be found in the [Trickest Workfl
 
 ### Continuous Integration 
 
-You can find the Github Action for the `trickest-cli` at https://github.com/trickest/action and the Docker image at https://quay.io/trickest/trickest-cli.
+You can find the GitHub Action for the `trickest-cli` at https://github.com/trickest/action and the Docker image at https://quay.io/trickest/trickest-cli.
 
 The `execute` command can be used as part of a CI pipeline to execute your Trickest workflows whenever your code or infrastructure changes. Optionally, you can use the `--watch` command inside the action to watch a workflow's progress until it completes. 
 
@@ -216,16 +216,17 @@ Use the **stop** command to stop a running workflow or node
 trickest stop --workflow <workflow_name> --space <space_name> [--run <run_id>] [--all] [--nodes <node_name_or_id>] [--child <child_task_index>]
 ```
 
-| Flag       | Type    | Default | Description                                                                                                                     |
-| ---------- | ------- | ------- | ------------------------------------------------------------------------------------------------------------------------------- |
-| --url      | string  | /       | URL copied from the Trickest platform, referencing a workflow and, optionally, a run/node                                       |
-| --workflow | string  | /       | The name of the workflow.                                                                                                       |
-| --project | string   | /       | The name of the project to which workflow belongs                                                                               |
-| --space    | string  | /       | The name of the space to which workflow belongs                                                                                 |
-| --run      | string  | /       | Stop a specific run                                                                                                             |
-| --all      | bool    | false   | Stop all runs                                                                                                                   |
-| --nodes    | string  | /       | A comma-separated list of nodes to stop. If none specified, the entire run will be stopped. If a node is a task group, the `--child` flag must be used               |
-| --child    | string  | /       | A comma-separated list or range of child tasks to stop. Example: `--child 1,2,3` or `--child 1-3` will stop the first three tasks in the specified node's taskgroup |
+| Flag       | Type   | Default | Description                                                                                                                                                         |
+|------------|--------|---------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| --url      | string | /       | URL copied from the Trickest platform, referencing a workflow and, optionally, a run/node                                                                           |
+| --workflow | string | /       | The name of the workflow.                                                                                                                                           |
+| --project  | string | /       | The name of the project to which workflow belongs                                                                                                                   |
+| --space    | string | /       | The name of the space to which workflow belongs                                                                                                                     |
+| --run      | string | /       | Stop a specific run                                                                                                                                                 |
+| --all      | bool   | false   | Stop all runs                                                                                                                                                       |
+| --nodes    | string | /       | A comma-separated list of nodes to stop. If none specified, the entire run will be stopped. If a node is a task group, the `--child` flag must be used              |
+| --child    | string | /       | A comma-separated list or range of child tasks to stop. Example: `--child 1,2,3` or `--child 1-3` will stop the first three tasks in the specified node's taskgroup |
+| --status   | string | running | A comma-separated list of run statuses to evaluate for stopping. Example: `pending,submitted,running`                                                               |
 
 ## Output
 Use the **output** command to download the outputs of your particular workflow execution(s) to your local environment.
@@ -442,7 +443,7 @@ trickest scripts delete [--id <script_id>] [--name <script_name>]
 ## Report Bugs / Feedback
 We look forward to any feedback you want to share with us or if you're stuck with a problem you can contact us at [[email protected]](mailto:[email protected]).
 
-You can also create an [Issue](https://github.com/trickest/trickest-cli/issues/new/choose) in the Github repository.
+You can also create an [Issue](https://github.com/trickest/trickest-cli/issues/new/choose) in the GitHub repository.
 
 [<img src="./banner.png" />](https://trickest.io/auth/register)
 

cmd/get/json_output.go

@@ -64,7 +64,7 @@ type JSONSubJob struct {
 func NewJSONRun(run *trickest.Run, subjobs []trickest.SubJob, taskGroupStatsMap map[uuid.UUID]stats.TaskGroupStats) *JSONRun {
 	jsonRun := &JSONRun{
 		ID:                  *run.ID,
-		Status:              run.Status,
+		Status:              run.Status.String(),
 		Author:              run.Author,
 		CreationType:        run.CreationType,
 		CreatedDate:         run.CreatedDate,

cmd/stop/stop.go

@@ -22,15 +22,18 @@ type Config struct {
 	Nodes          []string
 	ChildrenRanges []string
 	Children       []int
+	RunStatuses    []string
 }
 
 var cfg = &Config{}
+var defaultStatuses = trickest.RunStatuses{trickest.RunStatusRunning}.Strings()
 
 func init() {
 	StopCmd.Flags().BoolVar(&cfg.RunSpec.AllRuns, "all", false, "Stop all runs")
 	StopCmd.Flags().StringVar(&cfg.RunSpec.RunID, "run", "", "Stop a specific run")
 	StopCmd.Flags().StringSliceVar(&cfg.Nodes, "nodes", []string{}, "Nodes to stop. If none specified, the entire run will be stopped. If a node is a task group, the `--child` flag must be used (can be used multiple times)")
 	StopCmd.Flags().StringSliceVar(&cfg.ChildrenRanges, "child", []string{}, "Child tasks to stop. If a node is a task group, the `--child` flag must be used (can be used multiple times)")
+	StopCmd.Flags().StringSliceVar(&cfg.RunStatuses, "status", defaultStatuses, "Run statuses to be evaluated for stopping (can be used multiple times or comma-separated, e.g. `pending,submitted`)")
 }
 
 // StopCmd represents the stop command
@@ -93,14 +96,29 @@ func run(cfg *Config) error {
 		return fmt.Errorf("failed to create client: %w", err)
 	}
 
+	for _, status := range cfg.RunStatuses {
+		parsedStatus := trickest.RunStatus(strings.ToUpper(strings.TrimSpace(status)))
+		if !parsedStatus.Valid() {
+			return fmt.Errorf("invalid run status: %s", status)
+		}
+		if !parsedStatus.CanStop() {
+			return fmt.Errorf("cannot stop runs with status: %s", status)
+		}
+		cfg.RunSpec.RunStatuses = append(cfg.RunSpec.RunStatuses, parsedStatus)
+	}
+
 	ctx := context.Background()
 
-	cfg.RunSpec.RunStatus = "RUNNING"
 	runs, err := cfg.RunSpec.GetRuns(ctx, client)
 	if err != nil {
 		return fmt.Errorf("failed to get runs: %w", err)
 	}
 
+	if len(runs) == 0 {
+		fmt.Println("No runs found")
+		return nil
+	}
+
 	var errs []error
 	for _, run := range runs {
 		if len(cfg.Nodes) == 0 && len(cfg.Children) == 0 {

pkg/actions/output.go

@@ -46,7 +46,7 @@ func PrintDownloadResults(results []DownloadResult, runID uuid.UUID, destination
 // DownloadRunOutput downloads the outputs for the specified nodes in the run
 // Returns the download result summary, the directory where the outputs were saved, and an error if _all_ of the downloads failed
 func DownloadRunOutput(client *trickest.Client, run *trickest.Run, nodes []string, files []string, destinationPath string) ([]DownloadResult, string, error) {
-	if run.Status == "PENDING" || run.Status == "SUBMITTED" {
+	if !run.Status.IsStarted() {
 		return nil, "", fmt.Errorf("run %s has not started yet (status: %s)", run.ID.String(), run.Status)
 	}
 

pkg/config/workflowrunspec.go

@@ -15,7 +15,7 @@ type WorkflowRunSpec struct {
 	RunID        string
 	NumberOfRuns int
 	AllRuns      bool
-	RunStatus    string
+	RunStatuses  trickest.RunStatuses
 
 	// Workflow identification
 	SpaceName    string
@@ -36,8 +36,8 @@ func (s WorkflowRunSpec) GetRuns(ctx context.Context, client *trickest.Client) (
 		if err != nil {
 			return nil, err
 		}
-		if s.RunStatus != "" && run.Status != s.RunStatus {
-			return nil, fmt.Errorf("run %s has status %q, expected status %q", run.ID, run.Status, s.RunStatus)
+		if len(s.RunStatuses) > 0 && !s.RunStatuses.Contains(run.Status) {
+			return nil, fmt.Errorf("run %s has status %q, expected status %q", run.ID, run.Status, s.RunStatuses.String())
 		}
 		return []trickest.Run{*run}, nil
 	}
@@ -53,7 +53,7 @@ func (s WorkflowRunSpec) GetRuns(ctx context.Context, client *trickest.Client) (
 		limit = s.NumberOfRuns
 	}
 
-	runs, err := client.GetRuns(ctx, workflow.ID, s.RunStatus, limit)
+	runs, err := client.GetRuns(ctx, workflow.ID, s.RunStatuses.String(), limit)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get runs: %w", err)
 	}

pkg/display/run/printer.go

@@ -57,7 +57,7 @@ func (p *RunPrinter) PrintAll(run *trickest.Run, subJobs []trickest.SubJob, vers
 
 	// Print basic run details
 	output.WriteString(p.formatKeyValue("Name", run.WorkflowName))
-	output.WriteString(p.formatKeyValue("Status", run.Status))
+	output.WriteString(p.formatKeyValue("Status", run.Status.String()))
 	output.WriteString(p.formatKeyValue("Machines", fmt.Sprintf("%d", run.Machines)))
 	output.WriteString(p.formatKeyValue("Parallelism", fmt.Sprintf("%d", run.Parallelism)))
 	output.WriteString(p.formatKeyValue("Fleet", run.FleetName))
@@ -68,7 +68,7 @@ func (p *RunPrinter) PrintAll(run *trickest.Run, subJobs []trickest.SubJob, vers
 		output.WriteString(p.formatKeyValue("Created",
 			run.CreatedDate.In(time.Local).Format(time.RFC1123)+" ("+FormatDuration(time.Since(*run.CreatedDate))+" ago)"))
 	}
-	if run.Status != "PENDING" && run.StartedDate != nil {
+	if run.Status.IsStarted() && run.StartedDate != nil {
 		output.WriteString(p.formatKeyValue("Started",
 			run.StartedDate.In(time.Local).Format(time.RFC1123)+" ("+FormatDuration(time.Since(*run.StartedDate))+" ago)"))
 	}
@@ -110,7 +110,7 @@ func (p *RunPrinter) PrintAll(run *trickest.Run, subJobs []trickest.SubJob, vers
 	// - If the run is still running, mark them as "pending"
 	// - If the run is finished, mark them as "stopped"
 	defaultSubJobStatus := "stopped"
-	if run.Status == "RUNNING" {
+	if run.Status == trickest.RunStatusRunning {
 		defaultSubJobStatus = "pending"
 	}
 	output.WriteString(p.formatSubJobTree(subJobs, version, defaultSubJobStatus, includeTaskGroupStats))

pkg/trickest/run.go

@@ -6,16 +6,32 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"slices"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
 )
 
+const (
+	RunStatusPending   RunStatus = "PENDING"
+	RunStatusSubmitted RunStatus = "SUBMITTED"
+	RunStatusRunning   RunStatus = "RUNNING"
+	RunStatusCompleted RunStatus = "COMPLETED"
+	RunStatusFailed    RunStatus = "FAILED"
+	RunStatusStopping  RunStatus = "STOPPING"
+	RunStatusStopped   RunStatus = "STOPPED"
+)
+
+type RunStatus string
+
+type RunStatuses []RunStatus
+
 // Run represents a workflow run
 type Run struct {
 	ID                  *uuid.UUID         `json:"id,omitempty"`
 	Name                string             `json:"name,omitempty"`
-	Status              string             `json:"status,omitempty"`
+	Status              RunStatus          `json:"status,omitempty"`
 	Machines            int                `json:"machines,omitempty"`
 	Parallelism         int                `json:"parallelism,omitempty"`
 	WorkflowVersionInfo *uuid.UUID         `json:"workflow_version_info,omitempty"`
@@ -73,6 +89,55 @@ type RunSubJobInsights struct {
 	Stopped   int `json:"stopped"`
 }
 
+func (s RunStatus) Valid() bool {
+	switch s {
+	case RunStatusPending, RunStatusSubmitted, RunStatusRunning,
+		RunStatusCompleted, RunStatusFailed,
+		RunStatusStopping, RunStatusStopped:
+		return true
+	default:
+		return false
+	}
+}
+
+func (s RunStatus) IsStarted() bool {
+	switch s {
+	case RunStatusRunning, RunStatusCompleted, RunStatusFailed, RunStatusStopping, RunStatusStopped:
+		return true
+	default:
+		return false
+	}
+}
+
+func (s RunStatus) CanStop() bool {
+	switch s {
+	case RunStatusPending, RunStatusSubmitted, RunStatusRunning:
+		return true
+	default:
+		return false
+	}
+}
+
+func (s RunStatus) String() string {
+	return string(s)
+}
+
+func (s RunStatuses) Contains(status RunStatus) bool {
+	return slices.Contains(s, status)
+}
+
+func (s RunStatuses) Strings() []string {
+	result := make([]string, len(s))
+	for i, status := range s {
+		result[i] = status.String()
+	}
+	return result
+}
+
+func (s RunStatuses) String() string {
+	return strings.Join(s.Strings(), ",")
+}
+
 // GetRun retrieves a run by ID
 func (c *Client) GetRun(ctx context.Context, id uuid.UUID) (*Run, error) {
 	var run Run