From f4fb52c23788da835d178cd2e7a34d4959d622fd Mon Sep 17 00:00:00 2001
From: Kumar Krishna <kumar.krishna@gmail.com>
Date: Tue, 15 Nov 2016 22:24:07 -0800
Subject: [PATCH] New api GetTaskStatus

---
 examples/client.go | 21 ++++++++++++++++++---
 realis.go          | 21 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/examples/client.go b/examples/client.go
index 6d511b5..f2ef842 100644
--- a/examples/client.go
+++ b/examples/client.go
@@ -235,7 +235,7 @@ func main() {
 			os.Exit(1)
 		}
 
-		fmt.Println("Number of live instances: ", len(live))
+		fmt.Printf("Live instances: %+v\n", live)
 		break
 	case "activeCount":
 		fmt.Println("Getting instance count")
@@ -266,7 +266,7 @@ func main() {
 		fmt.Println(resp.String())
 		break
 	case "update":
-		fmt.Println("Updating a job with with more RAM and to 3 instances")
+		fmt.Println("Updating a job with with more RAM and to 5 instances")
 		taskConfig, err := r.FetchTaskConfig(aurora.InstanceKey{job.JobKey(), 0})
 		if err != nil {
 			fmt.Println(err)
@@ -282,7 +282,7 @@ func main() {
 		}
 
 		jobUpdateKey := response.JobUpdateKey(resp)
-		monitor.JobUpdate(*jobUpdateKey, 5, 100)
+		monitor.JobUpdate(*jobUpdateKey, 5, 500)
 
 		break
 	case "updateDetails":
@@ -316,12 +316,27 @@ func main() {
 	case "taskConfig":
 		fmt.Println("Getting job info")
 		config, err := r.FetchTaskConfig(aurora.InstanceKey{job.JobKey(), 0})
+
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
 		}
 		print(config.String())
 		break
+	case "taskStatus":
+		fmt.Println("Getting task status")
+		taskQ := &aurora.TaskQuery{Role: job.JobKey().Role,
+			Environment: job.JobKey().Environment,
+			JobName:     job.JobKey().Name,
+		}
+		tasks, err := r.GetTaskStatus(taskQ)
+		if err != nil {
+			fmt.Printf("error: %+v\n ",err )
+			os.Exit(1)
+		}
+		fmt.Printf("length: %d\n ", len(tasks))
+		fmt.Printf("tasks: %+v\n", tasks)
+
 	default:
 		fmt.Println("Command not supported")
 		os.Exit(1)
diff --git a/realis.go b/realis.go
index 3760a20..59f614f 100644
--- a/realis.go
+++ b/realis.go
@@ -33,6 +33,7 @@ type Realis interface {
 	AddInstances(instKey aurora.InstanceKey, count int32) (*aurora.Response, error)
 	CreateJob(auroraJob Job) (*aurora.Response, error)
 	DescheduleCronJob(key *aurora.JobKey) (*aurora.Response, error)
+	GetTaskStatus(query *aurora.TaskQuery)([]*aurora.ScheduledTask, error)
 	FetchTaskConfig(instKey aurora.InstanceKey) (*aurora.TaskConfig, error)
 	GetInstanceIds(key *aurora.JobKey, states map[aurora.ScheduleStatus]bool) (map[int32]bool, error)
 	JobUpdateDetails(updateQuery aurora.JobUpdateQuery) (*aurora.Response, error)
@@ -286,6 +287,21 @@ func (r realisClient) AddInstances(instKey aurora.InstanceKey, count int32) (*au
 	return response.ResponseCodeCheck(resp)
 }
 
+func (r realisClient) GetTaskStatus(query *aurora.TaskQuery)(tasks []*aurora.ScheduledTask, e error) {
+
+	resp, err := r.client.GetTasksStatus(query)
+	if err != nil {
+		return nil, errors.Wrap(err, "Error querying Aurora Scheduler for task status")
+	}
+	//Check for response code..
+	if resp.GetResponseCode() != aurora.ResponseCode_OK {
+		return nil, errors.New(resp.ResponseCode.String() + "--" + response.CombineMessage(resp))
+	}
+
+	return response.ScheduleStatusResult(resp).GetTasks(), nil
+}
+
+
 func (r realisClient) FetchTaskConfig(instKey aurora.InstanceKey) (*aurora.TaskConfig, error) {
 
 	ids := make(map[int32]bool)
@@ -302,6 +318,11 @@ func (r realisClient) FetchTaskConfig(instKey aurora.InstanceKey) (*aurora.TaskC
 		return nil, errors.Wrap(err, "Error querying Aurora Scheduler for task configuration")
 	}
 
+	//Check for response code..
+	if resp.GetResponseCode() != aurora.ResponseCode_OK {
+		return nil, errors.New(resp.ResponseCode.String() + "--" +response.CombineMessage(resp))
+	}
+
 	tasks := response.ScheduleStatusResult(resp).GetTasks()
 
 	if len(tasks) == 0 {