diff --git a/realis.go b/realis.go index 84c8b1e..62a9236 100644 --- a/realis.go +++ b/realis.go @@ -51,6 +51,7 @@ type Realis interface { GetTaskStatus(query *aurora.TaskQuery) ([]*aurora.ScheduledTask, error) GetTasksWithoutConfigs(query *aurora.TaskQuery) ([]*aurora.ScheduledTask, error) GetJobs(role string) (*aurora.Response, *aurora.GetJobsResult_, error) + GetPendingReason(query *aurora.TaskQuery) (pendingReasons []*aurora.PendingReason, e error) JobUpdateDetails(updateQuery aurora.JobUpdateQuery) (*aurora.Response, error) KillJob(key *aurora.JobKey) (*aurora.Response, error) KillInstances(key *aurora.JobKey, instances ...int32) (*aurora.Response, error) @@ -72,6 +73,7 @@ type Realis interface { // Admin functions DrainHosts(hosts ...string) (*aurora.Response, *aurora.DrainHostsResult_, error) + StartMaintenance(hosts ...string) (*aurora.Response, *aurora.StartMaintenanceResult_, error) EndMaintenance(hosts ...string) (*aurora.Response, *aurora.EndMaintenanceResult_, error) MaintenanceStatus(hosts ...string) (*aurora.Response, *aurora.MaintenanceStatusResult_, error) SetQuota(role string, cpu *float64, ram *int64, disk *int64) (*aurora.Response, error) @@ -862,6 +864,30 @@ func (r *realisClient) GetTaskStatus(query *aurora.TaskQuery) (tasks []*aurora.S return response.ScheduleStatusResult(resp).GetTasks(), nil } +// Get pending reason +func (r *realisClient) GetPendingReason(query *aurora.TaskQuery) (pendingReasons []*aurora.PendingReason, e error) { + + r.logger.DebugPrintf("GetPendingReason Thrift Payload: %+v\n", query) + + resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { + return r.client.GetPendingReason(query) + }) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Error querying Aurora Scheduler for pending Reasons") + } + + var result map[*aurora.PendingReason]bool + + if resp != nil && resp.GetResult_() != nil { + result = resp.GetResult_().GetGetPendingReasonResult_().GetReasons() + } + for reason := range result { + pendingReasons = append(pendingReasons, reason) + } + return pendingReasons, nil +} + // Get information about task including without a task configuration object func (r *realisClient) GetTasksWithoutConfigs(query *aurora.TaskQuery) (tasks []*aurora.ScheduledTask, e error) { @@ -984,6 +1010,37 @@ func (r *realisClient) DrainHosts(hosts ...string) (*aurora.Response, *aurora.Dr return resp, result, nil } +func (r *realisClient) StartMaintenance(hosts ...string) (*aurora.Response, *aurora.StartMaintenanceResult_, error) { + + var result *aurora.StartMaintenanceResult_ + + if len(hosts) == 0 { + return nil, nil, errors.New("no hosts provided to start maintenance on") + } + + hostList := aurora.NewHosts() + hostList.HostNames = make(map[string]bool) + for _, host := range hosts { + hostList.HostNames[host] = true + } + + r.logger.DebugPrintf("StartMaintenance Thrift Payload: %v\n", hostList) + + resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { + return r.adminClient.StartMaintenance(hostList) + }) + + if resp.GetResult_() != nil { + result = resp.GetResult_().GetStartMaintenanceResult_() + } + + if retryErr != nil { + return resp, result, errors.Wrap(retryErr, "Unable to recover connection") + } + + return resp, result, nil +} + func (r *realisClient) EndMaintenance(hosts ...string) (*aurora.Response, *aurora.EndMaintenanceResult_, error) { var result *aurora.EndMaintenanceResult_ @@ -1015,6 +1072,7 @@ func (r *realisClient) EndMaintenance(hosts ...string) (*aurora.Response, *auror return resp, result, nil } + func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *aurora.MaintenanceStatusResult_, error) { var result *aurora.MaintenanceStatusResult_