From 464db4e5ccc3d06bec537b4bd28cb905fd302f01 Mon Sep 17 00:00:00 2001
From: lawwong1 <105893156+lawwong1@users.noreply.github.com>
Date: Thu, 22 Sep 2022 20:05:19 -0700
Subject: [PATCH] Implement Australis api to restart a task (#43)

---
 cmd/restart.go                  | 63 +++++++++++++++++++++++++++++++++
 docs/australis_restart_tasks.md | 45 +++++++++++++++++++++++
 2 files changed, 108 insertions(+)
 create mode 100644 docs/australis_restart_tasks.md

diff --git a/cmd/restart.go b/cmd/restart.go
index d90d5fd..c207ffb 100644
--- a/cmd/restart.go
+++ b/cmd/restart.go
@@ -15,6 +15,11 @@
 package cmd
 
 import (
+	"strconv"
+	"strings"
+
+	realis "github.com/aurora-scheduler/gorealis/v2"
+
 	"github.com/aurora-scheduler/gorealis/v2/gen-go/apache/aurora"
 	"github.com/spf13/cobra"
 )
@@ -26,6 +31,17 @@ func init() {
 	restartJobCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment")
 	restartJobCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role")
 	restartJobCmd.Flags().StringVarP(name, "name", "n", "", "Aurora Name")
+
+	restartCmd.AddCommand(restartTasksCmd)
+	restartTasksCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment")
+	restartTasksCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role")
+	restartTasksCmd.Flags().StringVarP(name, "name", "n", "", "Aurora Name")
+	restartTasksCmd.Flags().StringVarP(instances, "instances", "I", "", "Instances e.g. 1, 2, 5")
+	restartTasksCmd.Flags().BoolVarP(&monitor, "monitor", "m", true, "monitor the result after sending the command")
+	restartTasksCmd.MarkFlagRequired("environment")
+	restartTasksCmd.MarkFlagRequired("role")
+	restartTasksCmd.MarkFlagRequired("name")
+	restartTasksCmd.MarkFlagRequired("instances")
 }
 
 var restartCmd = &cobra.Command{
@@ -39,9 +55,56 @@ var restartJobCmd = &cobra.Command{
 	Run:   restartJob,
 }
 
+var restartTasksCmd = &cobra.Command{
+	Use:   "tasks",
+	Short: "Restart tasks for a Job.",
+	Run:   restartTasks,
+}
+
 func restartJob(cmd *cobra.Command, args []string) {
 	key := aurora.JobKey{Environment: *env, Role: *role, Name: *name}
 	if err := client.RestartJob(key); err != nil {
 		log.Fatal("unable to create Aurora job: ", err)
 	}
 }
+
+func restartTasks(cmd *cobra.Command, args []string) {
+	log.Infof("Restarts task [Env:%s Role:%s Name:%s Instance:%s Monitor:%s]\n", *env, *role, *name, *instances, strconv.FormatBool(monitor))
+
+	//Set jobKey for the tasks to be killed.
+	task := realis.NewTask().
+		Environment(*env).
+		Role(*role).
+		Name(*name)
+
+	/*
+	* In the following block, we convert instance numbers, which were passed as strings, to integer values
+	* After converting them to integers, we add them to a slice of type int32.
+	 */
+
+	splitString := strings.Split(*instances, ",")
+	instanceList := make([]int32, len(splitString))
+
+	for i := range instanceList {
+		splitString[i] = strings.TrimSpace(splitString[i])
+		var instanceNumber int
+		var err error
+		if instanceNumber, err = strconv.Atoi(splitString[i]); err != nil {
+			log.Fatalln("Instance passed should be a number. Error: " + err.Error())
+			return
+		}
+		instanceList[i] = int32(instanceNumber)
+	}
+
+	//Call the RestartInstances function, passing the instanceList as the list of instances to be restarted.
+	if err := client.RestartInstances(task.JobKey(), instanceList...); err != nil {
+		log.Fatalln(err)
+	}
+
+	if monitor {
+		if ok, err := client.MonitorInstances(task.JobKey(), int32(len(instanceList)), 5, 50); !ok || err != nil {
+			log.Fatalln("Monitor failed to monitor the given task after restart. Error: " + err.Error())
+		}
+	}
+
+}
diff --git a/docs/australis_restart_tasks.md b/docs/australis_restart_tasks.md
new file mode 100644
index 0000000..1f33c79
--- /dev/null
+++ b/docs/australis_restart_tasks.md
@@ -0,0 +1,45 @@
+## australis restart tasks
+
+Restart tasks for a Job.
+
+### Synopsis
+
+Restart tasks for a Job.
+
+```
+australis restart tasks [flags]
+```
+
+### Options
+
+```
+  -e, --environment string   Aurora Environment
+  -h, --help                 help for tasks
+  -I, --instances string     Instances e.g. 1, 2, 5
+  -m, --monitor              monitor the result after sending the command (default true)
+  -n, --name string          Aurora Name
+  -r, --role string          Aurora Role
+```
+
+### Options inherited from parent commands
+
+```
+  -a, --caCertsPath string      Path where CA certificates can be found.
+  -c, --clientCert string       Client certificate to use to connect to Aurora.
+  -k, --clientKey string        Client key to use to connect to Aurora.
+      --config string           Config file to use. (default "/etc/aurora/australis.yml")
+  -l, --logLevel string         Set logging level [panic fatal error warning info debug trace]. (default "info")
+  -p, --password string         Password to use for API authentication
+  -s, --scheduler_addr string   Aurora Scheduler's address.
+  -i, --skipCertVerification    Skip CA certificate hostname verification.
+  -t, --timeout duration        Gorealis timeout. (default 20s)
+      --toJSON                  Print output in JSON format.
+  -u, --username string         Username to use for API authentication
+  -z, --zookeeper string        Zookeeper node(s) where Aurora stores information. (comma separated list)
+```
+
+### SEE ALSO
+
+* [australis restart](australis_restart.md)	 - Restart an Aurora Job.
+
+###### Auto generated by spf13/cobra on 21-Sep-2022