diff --git a/realis.go b/realis.go index 6c2f6d6..dde0e6d 100644 --- a/realis.go +++ b/realis.go @@ -71,6 +71,8 @@ type Realis interface { DrainHosts(hosts ...string) (*aurora.Response, *aurora.DrainHostsResult_, 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) + GetQuota(role string) (*aurora.Response, error) } type realisClient struct { @@ -944,3 +946,44 @@ func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *au return resp, result, nil } + +// SetQuota sets a quota aggregate for the given role +// TODO(zircote) Currently investigating an error that is returned from thrift calls that include resources for `NamedPort` and `NumGpu` +func (r *realisClient) SetQuota(role string, cpu *float64, ramMb *int64, diskMb *int64) (*aurora.Response, error) { + ram := aurora.NewResource() + ram.RamMb = ramMb + c := aurora.NewResource() + c.NumCpus = cpu + d := aurora.NewResource() + d.DiskMb = diskMb + quota := aurora.NewResourceAggregate() + quota.Resources = make(map[*aurora.Resource]bool) + quota.Resources[ram] = true + quota.Resources[c] = true + quota.Resources[d] = true + resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { + resp, retryErr := r.adminClient.SetQuota(role, quota) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to set role quota") + } + return resp, nil + }) + return resp, retryErr + +} + +// GetQuota returns the resource aggregate for the given role +func (r *realisClient) GetQuota(role string) (*aurora.Response, error) { + + resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { + resp, retryErr :=r.adminClient.GetQuota(role) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to get role quota") + } + return resp, nil + }) + + return resp, retryErr +} diff --git a/realis_e2e_test.go b/realis_e2e_test.go index 34817ab..ae61cf7 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -389,3 +389,37 @@ func TestRealisClient_SessionThreadSafety(t *testing.T) { wg.Wait() } + +// Test setting and getting the quota +func TestRealisClient_SetQuota(t *testing.T) { + var cpu = 3.5 + var ram int64 = 20480 + var disk int64 = 10240 + resp, err := r.SetQuota("vagrant", &cpu, &ram, &disk) + assert.NoError(t, err) + assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) + t.Run("TestRealisClient_GetQuota", func(t *testing.T) { + // Test GetQuota based on previously set values + var result *aurora.GetQuotaResult_ + resp, err = r.GetQuota("vagrant") + if resp.GetResult_() != nil { + result = resp.GetResult_().GetQuotaResult_ + } + assert.NoError(t, err) + assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) + for res := range result.Quota.GetResources() { + switch true { + case res.DiskMb != nil: + assert.Equal(t, disk, *res.DiskMb) + break + case res.NumCpus != nil: + assert.Equal(t, cpu, *res.NumCpus) + break + case res.RamMb != nil: + assert.Equal(t, ram, *res.RamMb) + break + } + } + fmt.Print("GetQuota Result", result.String()) + }) +} diff --git a/updatejob.go b/updatejob.go index a1be06a..4ee4f14 100644 --- a/updatejob.go +++ b/updatejob.go @@ -138,7 +138,6 @@ func (u *UpdateJob) RollbackOnFail(rollback bool) *UpdateJob { return u } - func NewUpdateSettings() *aurora.JobUpdateSettings { us := new(aurora.JobUpdateSettings)