diff --git a/realis.go b/realis.go index 6c2f6d6..a1512c7 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,34 @@ func (r *realisClient) MaintenanceStatus(hosts ...string) (*aurora.Response, *au return resp, result, nil } + +// SetQuota sets a quota aggregate for the given role +func (r *realisClient) SetQuota(role string, cpu float64, ramMb int64, diskMb int64) (*aurora.Response, error) { + quota := &aurora.ResourceAggregate{ + NumCpus: cpu, + RamMb: ramMb, + DiskMb: diskMb, + } + resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) { + return r.adminClient.SetQuota(role, quota) + }) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to set role quota") + } + return resp, nil + +} + +// 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) { + return r.adminClient.GetQuota(role) + }) + + if retryErr != nil { + return nil, errors.Wrap(retryErr, "Unable to get role quota") + } + return resp, nil +} diff --git a/realis_e2e_test.go b/realis_e2e_test.go index 34817ab..1528594 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -389,3 +389,26 @@ func TestRealisClient_SessionThreadSafety(t *testing.T) { wg.Wait() } + +// Test setting and getting the quota +func TestRealisClient_SetQuota(t *testing.T) { + var cpu = 2.5 + var ram int64 = 10240 + var disk int64 = 10240 + resp, err := r.SetQuota("vagrant", cpu, ram, disk) + assert.NoError(t, err) + assert.Equal(t, aurora.ResponseCode_OK, resp.ResponseCode) + + // 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) + assert.Equal(t, cpu, result.Quota.NumCpus) + assert.Equal(t, ram, result.Quota.RamMb) + assert.Equal(t, disk, result.Quota.DiskMb) + 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)