From f88a63a40c99c798c4e81f36620fe233259e700a Mon Sep 17 00:00:00 2001
From: Samir Alajmovic <5246600+alajmo@users.noreply.github.com>
Date: Sat, 30 Jul 2022 22:49:31 +0200
Subject: [PATCH] Fix correct exit code on remote/local task error (#27)

- Allow duplicate hosts
- Fix correct exit code on remote/local task errors
- Fix local WorkDir when it's not explicitly set
---
 Makefile                                 |  2 +-
 core/errors.go                           | 14 +++++++++-
 core/run/exec.go                         | 11 +++++---
 core/run/table.go                        | 22 ++++++++++-----
 core/run/text.go                         | 34 +++++++++++++++---------
 core/run/unix.go                         |  2 +-
 core/run/windows.go                      |  1 +
 core/sake.1                              |  2 +-
 docs/changelog.md                        |  5 +++-
 test/Dockerfile                          |  2 +-
 test/integration/golden/golden-15.stdout |  3 ++-
 test/integration/golden/golden-18.stdout |  3 ++-
 test/integration/run_test.go             |  4 +--
 test/playground/sake.yaml                |  4 +++
 test/sake.yaml                           |  9 ++++++-
 test/{sake-1.yaml => servers.yaml}       |  2 +-
 test/{sake-2.yaml => tasks.yaml}         |  0
 17 files changed, 86 insertions(+), 34 deletions(-)
 rename test/{sake-1.yaml => servers.yaml} (98%)
 rename test/{sake-2.yaml => tasks.yaml} (100%)

diff --git a/Makefile b/Makefile
index d833fa7..916992f 100644
--- a/Makefile
+++ b/Makefile
@@ -2,7 +2,7 @@ NAME    := sake
 PACKAGE := github.com/alajmo/$(NAME)
 DATE    := $(shell date +%FT%T%Z)
 GIT     := $(shell [ -d .git ] && git rev-parse --short HEAD)
-VERSION := v0.10.1
+VERSION := v0.10.2
 
 default: build
 
diff --git a/core/errors.go b/core/errors.go
index 815022b..12aa1c4 100644
--- a/core/errors.go
+++ b/core/errors.go
@@ -152,6 +152,15 @@ func (f *TemplateParseError) Error() string {
 	return fmt.Sprintf("failed to parse %s", f.Msg)
 }
 
+type ExecError struct {
+	Err      error
+	ExitCode int
+}
+
+func (e *ExecError) Error() string {
+	return ""
+}
+
 func CheckIfError(err error) {
 	if err != nil {
 		Exit(err)
@@ -159,11 +168,14 @@ func CheckIfError(err error) {
 }
 
 func Exit(err error) {
-	switch err.(type) {
+	switch err := err.(type) {
 	case *ConfigErr:
 		// Errors are already mapped with `error:` prefix
 		fmt.Fprintf(os.Stderr, "%v", err)
 		os.Exit(1)
+	case *ExecError:
+		// Don't print anything when there's a ExecError: server execution failed
+		os.Exit(err.ExitCode)
 	default:
 		fmt.Fprintf(os.Stderr, "%s: %v\n", text.FgRed.Sprintf("error"), err)
 		os.Exit(1)
diff --git a/core/run/exec.go b/core/run/exec.go
index 5a77dde..21f6e38 100644
--- a/core/run/exec.go
+++ b/core/run/exec.go
@@ -90,7 +90,7 @@ func (run *Run) RunTask(
 		print.PrintTable("Unreachable Hosts", unreachableOutput.Rows, options, unreachableOutput.Headers[0:1], unreachableOutput.Headers[1:])
 
 		if !task.Spec.IgnoreUnreachable {
-			return nil
+			return &core.ExecError{Err: err, ExitCode: 4}
 		}
 	}
 
@@ -119,16 +119,19 @@ func (run *Run) RunTask(
 		spinner := core.GetSpinner()
 		spinner.Start(" Running", 500)
 
-		data := run.Table(runFlags.DryRun)
+		data, err := run.Table(runFlags.DryRun)
 		options := print.PrintTableOptions{Theme: task.Theme, OmitEmpty: task.Spec.OmitEmpty, Output: task.Spec.Output, SuppressEmptyColumns: false}
 		run.CleanupClients()
-
 		spinner.Stop()
-
 		print.PrintTable("", data.Rows, options, data.Headers[0:1], data.Headers[1:])
+
+		if err != nil {
+			return err
+		}
 	default:
 		err := run.Text(runFlags.DryRun)
 		run.CleanupClients()
+
 		if err != nil {
 			return err
 		}
diff --git a/core/run/table.go b/core/run/table.go
index b7a78e6..99894c4 100644
--- a/core/run/table.go
+++ b/core/run/table.go
@@ -2,9 +2,11 @@ package run
 
 import (
 	"fmt"
+	"golang.org/x/crypto/ssh"
 	"io"
 	"io/ioutil"
 	"os"
+	"os/exec"
 	"strings"
 	"sync"
 
@@ -12,7 +14,7 @@ import (
 	"github.com/alajmo/sake/core/dao"
 )
 
-func (run *Run) Table(dryRun bool) dao.TableOutput {
+func (run *Run) Table(dryRun bool) (dao.TableOutput, error) {
 	task := run.Task
 	servers := run.Servers
 
@@ -51,7 +53,7 @@ func (run *Run) Table(dryRun bool) dao.TableOutput {
 		if task.Spec.Parallel {
 			go func(i int, wg *sync.WaitGroup) {
 				defer wg.Done()
-				// TODO
+				// TODO: Handle errors when running tasks in parallel
 				_ = run.TableWork(i, dryRun, data, &dataMutex)
 			}(i, &wg)
 		} else {
@@ -62,15 +64,23 @@ func (run *Run) Table(dryRun bool) dao.TableOutput {
 				return err
 			}(i, &wg)
 
-			if run.Task.Spec.AnyErrorsFatal && err != nil {
-				break
+			if err != nil && run.Task.Spec.AnyErrorsFatal {
+				// Return proper exit code for failed tasks
+				switch err := err.(type) {
+				case *ssh.ExitError:
+					return data, &core.ExecError{Err: err, ExitCode: err.ExitStatus()}
+				case *exec.ExitError:
+					return data, &core.ExecError{Err: err, ExitCode: err.ExitCode()}
+				default:
+					return data, err
+				}
 			}
 		}
 
 	}
 	wg.Wait()
 
-	return data
+	return data, nil
 }
 
 func (run *Run) TableWork(rIndex int, dryRun bool, data dao.TableOutput, dataMutex *sync.RWMutex) error {
@@ -104,7 +114,7 @@ func (run *Run) TableWork(rIndex int, dryRun bool, data dao.TableOutput, dataMut
 		}
 
 		err := RunTableCmd(tableCmd, data, dataMutex, &wg)
-		if err != nil && !task.Spec.IgnoreErrors {
+		if !task.Spec.IgnoreErrors && err != nil {
 			return err
 		}
 	}
diff --git a/core/run/text.go b/core/run/text.go
index 254f93c..bf3563f 100644
--- a/core/run/text.go
+++ b/core/run/text.go
@@ -5,10 +5,12 @@ import (
 	"bytes"
 	"fmt"
 	"github.com/jedib0t/go-pretty/v6/text"
+	"golang.org/x/crypto/ssh"
 	"golang.org/x/exp/slices"
 	"golang.org/x/term"
 	"io"
 	"os"
+	"os/exec"
 	"strings"
 	"sync"
 	"text/template"
@@ -29,7 +31,7 @@ func (run *Run) Text(dryRun bool) error {
 		if run.Task.Spec.Parallel {
 			go func(i int, wg *sync.WaitGroup) {
 				defer wg.Done()
-				// TODO
+				// TODO: Handle errors when running tasks in parallel
 				_ = run.TextWork(i, prefixMaxLen, dryRun)
 			}(i, &wg)
 		} else {
@@ -39,16 +41,24 @@ func (run *Run) Text(dryRun bool) error {
 				return err
 			}(i, &wg)
 
-			switch err.(type) {
-			case *template.ExecError:
-				return err
-			case *core.TemplateParseError:
-				return err
-			default:
-				if run.Task.Spec.AnyErrorsFatal && err != nil {
-					// The error is printed for each server in method RunTextCmd.
-					// We just return early so other tasks are not executed.
-					return nil
+			if err != nil {
+				switch err.(type) {
+				case *template.ExecError:
+					return err
+				case *core.TemplateParseError:
+					return err
+				default:
+					if run.Task.Spec.AnyErrorsFatal {
+						// Return proper exit code for failed tasks
+						switch err := err.(type) {
+						case *ssh.ExitError:
+							return &core.ExecError{Err: err, ExitCode: err.ExitStatus()}
+						case *exec.ExitError:
+							return &core.ExecError{Err: err, ExitCode: err.ExitCode()}
+						default:
+							return err
+						}
+					}
 				}
 			}
 		}
@@ -102,7 +112,7 @@ func (run *Run) TextWork(rIndex int, prefixMaxLen int, dryRun bool) error {
 		case *core.TemplateParseError:
 			return err
 		default:
-			if err != nil && !task.Spec.IgnoreErrors {
+			if !task.Spec.IgnoreErrors && err != nil {
 				return err
 			}
 		}
diff --git a/core/run/unix.go b/core/run/unix.go
index 426c779..b589e0a 100644
--- a/core/run/unix.go
+++ b/core/run/unix.go
@@ -1,3 +1,4 @@
+//go:build !windows
 // +build !windows
 
 package run
@@ -46,4 +47,3 @@ func ExecTTY(cmd string, envs []string) error {
 
 	return nil
 }
-
diff --git a/core/run/windows.go b/core/run/windows.go
index 42aef3b..3e80295 100644
--- a/core/run/windows.go
+++ b/core/run/windows.go
@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package run
diff --git a/core/sake.1 b/core/sake.1
index 856a89f..3f8d270 100644
--- a/core/sake.1
+++ b/core/sake.1
@@ -1,4 +1,4 @@
-.TH "SAKE" "1" "2022-06-26T09:26:27CEST" "v0.10.0" "Sake Manual" "sake"
+.TH "SAKE" "1" "2022-07-30T22:44:33CEST" "v0.10.2" "Sake Manual" "sake"
 .SH NAME
 sake - sake is a command runner for local and remote hosts
 
diff --git a/docs/changelog.md b/docs/changelog.md
index 247ca1f..e01652c 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,10 +2,13 @@
 
 ## Unreleased
 
+## 0.10.2
+
 ### Fixes
 
-- Fix local WorkDir when it's not set
 - Allow duplicate hosts
+- Fix correct exit code on remote/local task errors (#27)
+- Fix local WorkDir when it's not explicitly set
 
 ## 0.10.1
 
diff --git a/test/Dockerfile b/test/Dockerfile
index 543ac3c..c1bf025 100644
--- a/test/Dockerfile
+++ b/test/Dockerfile
@@ -1,4 +1,4 @@
-FROM ubuntu:21.10
+FROM ubuntu:latest
 
 RUN apt-get update && apt-get install openssh-server sudo -y
 
diff --git a/test/integration/golden/golden-15.stdout b/test/integration/golden/golden-15.stdout
index dd92117..48013ba 100755
--- a/test/integration/golden/golden-15.stdout
+++ b/test/integration/golden/golden-15.stdout
@@ -1,7 +1,7 @@
 Index: 15
 Name: fatal true
 Cmd: go run ../../main.go run fatal-true -t reachable
-WantErr: false
+WantErr: true
 
 ---
 
@@ -19,3 +19,4 @@ WantErr: false
  server-8  |               
  server-9  |               
 
+exit status 1
diff --git a/test/integration/golden/golden-18.stdout b/test/integration/golden/golden-18.stdout
index 09821b2..c6adfc6 100755
--- a/test/integration/golden/golden-18.stdout
+++ b/test/integration/golden/golden-18.stdout
@@ -1,7 +1,7 @@
 Index: 18
 Name: unreachable false
 Cmd: go run ../../main.go run unreachable -a
-WantErr: false
+WantErr: true
 
 ---
 
@@ -11,3 +11,4 @@ Unreachable Hosts
 -------------+-----------------+------+------+------------------------------------------------
  unreachable | unreachable.lan | test | 22   | dial tcp: lookup unreachable.lan: no such host 
 
+exit status 4
diff --git a/test/integration/run_test.go b/test/integration/run_test.go
index af8d475..4f2ea01 100644
--- a/test/integration/run_test.go
+++ b/test/integration/run_test.go
@@ -95,7 +95,7 @@ var cases = []TemplateTest{
 	{
 		TestName:   "fatal true",
 		TestCmd: "go run ../../main.go run fatal-true -t reachable",
-		WantErr: false,
+		WantErr: true,
 	},
 	{
 		TestName:   "ignore_errors false",
@@ -110,7 +110,7 @@ var cases = []TemplateTest{
 	{
 		TestName:   "unreachable false",
 		TestCmd: "go run ../../main.go run unreachable -a",
-		WantErr: false,
+		WantErr: true,
 	},
 	{
 		TestName:   "unreachable true",
diff --git a/test/playground/sake.yaml b/test/playground/sake.yaml
index 6b5f144..5a96f5c 100644
--- a/test/playground/sake.yaml
+++ b/test/playground/sake.yaml
@@ -58,6 +58,10 @@ env:
   DATE: $(date -u +"%Y-%m-%dT%H:%M:%S%Z")
 
 tasks:
+  exit:
+    local: true
+    cmd: exit 3
+
   ping:
     target: all
     desc: ping server
diff --git a/test/sake.yaml b/test/sake.yaml
index d8074e2..bc47026 100644
--- a/test/sake.yaml
+++ b/test/sake.yaml
@@ -1,7 +1,7 @@
 disable_verify_host: true
 
 import:
-  - ./sake-1.yaml
+  - ./servers.yaml
 
 env:
   NO_COLOR: true
@@ -24,6 +24,13 @@ specs:
       ignore_unreachable: true
       any_errors_fatal: false
 
+themes:
+    default:
+        text:
+            prefix: true
+            header: '{{ .Style "TASK" "bold" }}{{ if ne .NumTasks 1 }} ({{ .Index }}/{{ .NumTasks }}){{end}}{{ if and .Name .Desc }} [{{.Style .Name "bold"}}: {{ .Desc }}] {{ else if .Name }} [{{ .Name }}] {{ else if .Desc }} [{{ .Desc }}] {{end}}'
+        table:
+
 tasks:
   ping:
     target: all
diff --git a/test/sake-1.yaml b/test/servers.yaml
similarity index 98%
rename from test/sake-1.yaml
rename to test/servers.yaml
index 7b33a59..f389dc3 100644
--- a/test/sake-1.yaml
+++ b/test/servers.yaml
@@ -1,5 +1,5 @@
 import:
-  - ./sake-2.yaml
+  - ./tasks.yaml
 
 servers:
   localhost:
diff --git a/test/sake-2.yaml b/test/tasks.yaml
similarity index 100%
rename from test/sake-2.yaml
rename to test/tasks.yaml