mirror of https://github.com/grafana/grafana
Merge branch 'quotas' of https://github.com/raintank/grafana into raintank-quotas
Torkel Ödegaard
10 years ago
commit
6a30511fc4
@ -0,0 +1,68 @@ |
||||
package api |
||||
|
||||
import ( |
||||
"github.com/grafana/grafana/pkg/bus" |
||||
"github.com/grafana/grafana/pkg/middleware" |
||||
m "github.com/grafana/grafana/pkg/models" |
||||
"github.com/grafana/grafana/pkg/setting" |
||||
) |
||||
|
||||
func GetOrgQuotas(c *middleware.Context) Response { |
||||
if !setting.Quota.Enabled { |
||||
return ApiError(404, "Quotas not enabled", nil) |
||||
} |
||||
query := m.GetOrgQuotasQuery{OrgId: c.ParamsInt64(":orgId")} |
||||
|
||||
if err := bus.Dispatch(&query); err != nil { |
||||
return ApiError(500, "Failed to get org quotas", err) |
||||
} |
||||
|
||||
return Json(200, query.Result) |
||||
} |
||||
|
||||
func UpdateOrgQuota(c *middleware.Context, cmd m.UpdateOrgQuotaCmd) Response { |
||||
if !setting.Quota.Enabled { |
||||
return ApiError(404, "Quotas not enabled", nil) |
||||
} |
||||
cmd.OrgId = c.ParamsInt64(":orgId") |
||||
cmd.Target = c.Params(":target") |
||||
|
||||
if _, ok := m.QuotaToMap(setting.Quota.Org)[cmd.Target]; !ok { |
||||
return ApiError(404, "Invalid quota target", nil) |
||||
} |
||||
|
||||
if err := bus.Dispatch(&cmd); err != nil { |
||||
return ApiError(500, "Failed to update org quotas", err) |
||||
} |
||||
return ApiSuccess("Organization quota updated") |
||||
} |
||||
|
||||
func GetUserQuotas(c *middleware.Context) Response { |
||||
if !setting.Quota.Enabled { |
||||
return ApiError(404, "Quotas not enabled", nil) |
||||
} |
||||
query := m.GetUserQuotasQuery{UserId: c.ParamsInt64(":id")} |
||||
|
||||
if err := bus.Dispatch(&query); err != nil { |
||||
return ApiError(500, "Failed to get org quotas", err) |
||||
} |
||||
|
||||
return Json(200, query.Result) |
||||
} |
||||
|
||||
func UpdateUserQuota(c *middleware.Context, cmd m.UpdateUserQuotaCmd) Response { |
||||
if !setting.Quota.Enabled { |
||||
return ApiError(404, "Quotas not enabled", nil) |
||||
} |
||||
cmd.UserId = c.ParamsInt64(":id") |
||||
cmd.Target = c.Params(":target") |
||||
|
||||
if _, ok := m.QuotaToMap(setting.Quota.User)[cmd.Target]; !ok { |
||||
return ApiError(404, "Invalid quota target", nil) |
||||
} |
||||
|
||||
if err := bus.Dispatch(&cmd); err != nil { |
||||
return ApiError(500, "Failed to update org quotas", err) |
||||
} |
||||
return ApiSuccess("Organization quota updated") |
||||
} |
||||
@ -0,0 +1,144 @@ |
||||
package middleware |
||||
|
||||
import ( |
||||
"github.com/grafana/grafana/pkg/bus" |
||||
m "github.com/grafana/grafana/pkg/models" |
||||
"github.com/grafana/grafana/pkg/setting" |
||||
. "github.com/smartystreets/goconvey/convey" |
||||
"testing" |
||||
) |
||||
|
||||
func TestMiddlewareQuota(t *testing.T) { |
||||
|
||||
Convey("Given the grafana quota middleware", t, func() { |
||||
setting.Quota = setting.QuotaSettings{ |
||||
Enabled: true, |
||||
Org: &setting.OrgQuota{ |
||||
User: 5, |
||||
Dashboard: 5, |
||||
DataSource: 5, |
||||
ApiKey: 5, |
||||
}, |
||||
User: &setting.UserQuota{ |
||||
Org: 5, |
||||
}, |
||||
Global: &setting.GlobalQuota{ |
||||
Org: 5, |
||||
User: 5, |
||||
Dashboard: 5, |
||||
DataSource: 5, |
||||
ApiKey: 5, |
||||
Session: 5, |
||||
}, |
||||
} |
||||
|
||||
middlewareScenario("with user not logged in", func(sc *scenarioContext) { |
||||
bus.AddHandler("globalQuota", func(query *m.GetGlobalQuotaByTargetQuery) error { |
||||
query.Result = &m.GlobalQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: query.Default, |
||||
Used: 4, |
||||
} |
||||
return nil |
||||
}) |
||||
Convey("global quota not reached", func() { |
||||
sc.m.Get("/user", Quota("user"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/user").exec() |
||||
So(sc.resp.Code, ShouldEqual, 200) |
||||
}) |
||||
Convey("global quota reached", func() { |
||||
setting.Quota.Global.User = 4 |
||||
sc.m.Get("/user", Quota("user"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/user").exec() |
||||
So(sc.resp.Code, ShouldEqual, 403) |
||||
}) |
||||
Convey("global session quota not reached", func() { |
||||
setting.Quota.Global.Session = 10 |
||||
sc.m.Get("/user", Quota("session"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/user").exec() |
||||
So(sc.resp.Code, ShouldEqual, 200) |
||||
}) |
||||
Convey("global session quota reached", func() { |
||||
setting.Quota.Global.Session = 1 |
||||
sc.m.Get("/user", Quota("session"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/user").exec() |
||||
So(sc.resp.Code, ShouldEqual, 403) |
||||
}) |
||||
}) |
||||
|
||||
middlewareScenario("with user logged in", func(sc *scenarioContext) { |
||||
// log us in, so we have a user_id and org_id in the context
|
||||
sc.fakeReq("GET", "/").handler(func(c *Context) { |
||||
c.Session.Set(SESS_KEY_USERID, int64(12)) |
||||
}).exec() |
||||
|
||||
bus.AddHandler("test", func(query *m.GetSignedInUserQuery) error { |
||||
query.Result = &m.SignedInUser{OrgId: 2, UserId: 12} |
||||
return nil |
||||
}) |
||||
bus.AddHandler("globalQuota", func(query *m.GetGlobalQuotaByTargetQuery) error { |
||||
query.Result = &m.GlobalQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: query.Default, |
||||
Used: 4, |
||||
} |
||||
return nil |
||||
}) |
||||
bus.AddHandler("userQuota", func(query *m.GetUserQuotaByTargetQuery) error { |
||||
query.Result = &m.UserQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: query.Default, |
||||
Used: 4, |
||||
} |
||||
return nil |
||||
}) |
||||
bus.AddHandler("orgQuota", func(query *m.GetOrgQuotaByTargetQuery) error { |
||||
query.Result = &m.OrgQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: query.Default, |
||||
Used: 4, |
||||
} |
||||
return nil |
||||
}) |
||||
Convey("global datasource quota reached", func() { |
||||
setting.Quota.Global.DataSource = 4 |
||||
sc.m.Get("/ds", Quota("data_source"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/ds").exec() |
||||
So(sc.resp.Code, ShouldEqual, 403) |
||||
}) |
||||
Convey("user Org quota not reached", func() { |
||||
setting.Quota.User.Org = 5 |
||||
sc.m.Get("/org", Quota("org"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/org").exec() |
||||
So(sc.resp.Code, ShouldEqual, 200) |
||||
}) |
||||
Convey("user Org quota reached", func() { |
||||
setting.Quota.User.Org = 4 |
||||
sc.m.Get("/org", Quota("org"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/org").exec() |
||||
So(sc.resp.Code, ShouldEqual, 403) |
||||
}) |
||||
Convey("org dashboard quota not reached", func() { |
||||
setting.Quota.Org.Dashboard = 10 |
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/dashboard").exec() |
||||
So(sc.resp.Code, ShouldEqual, 200) |
||||
}) |
||||
Convey("org dashboard quota reached", func() { |
||||
setting.Quota.Org.Dashboard = 4 |
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/dashboard").exec() |
||||
So(sc.resp.Code, ShouldEqual, 403) |
||||
}) |
||||
Convey("org dashboard quota reached but quotas disabled", func() { |
||||
setting.Quota.Org.Dashboard = 4 |
||||
setting.Quota.Enabled = false |
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler) |
||||
sc.fakeReq("GET", "/dashboard").exec() |
||||
So(sc.resp.Code, ShouldEqual, 200) |
||||
}) |
||||
|
||||
}) |
||||
|
||||
}) |
||||
} |
||||
@ -0,0 +1,130 @@ |
||||
package models |
||||
|
||||
import ( |
||||
"errors" |
||||
"github.com/grafana/grafana/pkg/setting" |
||||
"time" |
||||
) |
||||
|
||||
var ErrInvalidQuotaTarget = errors.New("Invalid quota target") |
||||
|
||||
type Quota struct { |
||||
Id int64 |
||||
OrgId int64 |
||||
UserId int64 |
||||
Target string |
||||
Limit int64 |
||||
Created time.Time |
||||
Updated time.Time |
||||
} |
||||
|
||||
type QuotaScope struct { |
||||
Name string |
||||
Target string |
||||
DefaultLimit int64 |
||||
} |
||||
|
||||
type OrgQuotaDTO struct { |
||||
OrgId int64 `json:"org_id"` |
||||
Target string `json:"target"` |
||||
Limit int64 `json:"limit"` |
||||
Used int64 `json:"used"` |
||||
} |
||||
|
||||
type UserQuotaDTO struct { |
||||
UserId int64 `json:"user_id"` |
||||
Target string `json:"target"` |
||||
Limit int64 `json:"limit"` |
||||
Used int64 `json:"used"` |
||||
} |
||||
|
||||
type GlobalQuotaDTO struct { |
||||
Target string `json:"target"` |
||||
Limit int64 `json:"limit"` |
||||
Used int64 `json:"used"` |
||||
} |
||||
|
||||
type GetOrgQuotaByTargetQuery struct { |
||||
Target string |
||||
OrgId int64 |
||||
Default int64 |
||||
Result *OrgQuotaDTO |
||||
} |
||||
|
||||
type GetOrgQuotasQuery struct { |
||||
OrgId int64 |
||||
Result []*OrgQuotaDTO |
||||
} |
||||
|
||||
type GetUserQuotaByTargetQuery struct { |
||||
Target string |
||||
UserId int64 |
||||
Default int64 |
||||
Result *UserQuotaDTO |
||||
} |
||||
|
||||
type GetUserQuotasQuery struct { |
||||
UserId int64 |
||||
Result []*UserQuotaDTO |
||||
} |
||||
|
||||
type GetGlobalQuotaByTargetQuery struct { |
||||
Target string |
||||
Default int64 |
||||
Result *GlobalQuotaDTO |
||||
} |
||||
|
||||
type UpdateOrgQuotaCmd struct { |
||||
Target string `json:"target"` |
||||
Limit int64 `json:"limit"` |
||||
OrgId int64 `json:"-"` |
||||
} |
||||
|
||||
type UpdateUserQuotaCmd struct { |
||||
Target string `json:"target"` |
||||
Limit int64 `json:"limit"` |
||||
UserId int64 `json:"-"` |
||||
} |
||||
|
||||
func GetQuotaScopes(target string) ([]QuotaScope, error) { |
||||
scopes := make([]QuotaScope, 0) |
||||
switch target { |
||||
case "user": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.User}, |
||||
QuotaScope{Name: "org", Target: "org_user", DefaultLimit: setting.Quota.Org.User}, |
||||
) |
||||
return scopes, nil |
||||
case "org": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.Org}, |
||||
QuotaScope{Name: "user", Target: "org_user", DefaultLimit: setting.Quota.User.Org}, |
||||
) |
||||
return scopes, nil |
||||
case "dashboard": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.Dashboard}, |
||||
QuotaScope{Name: "org", Target: target, DefaultLimit: setting.Quota.Org.Dashboard}, |
||||
) |
||||
return scopes, nil |
||||
case "data_source": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.DataSource}, |
||||
QuotaScope{Name: "org", Target: target, DefaultLimit: setting.Quota.Org.DataSource}, |
||||
) |
||||
return scopes, nil |
||||
case "api_key": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.ApiKey}, |
||||
QuotaScope{Name: "org", Target: target, DefaultLimit: setting.Quota.Org.ApiKey}, |
||||
) |
||||
return scopes, nil |
||||
case "session": |
||||
scopes = append(scopes, |
||||
QuotaScope{Name: "global", Target: target, DefaultLimit: setting.Quota.Global.Session}, |
||||
) |
||||
return scopes, nil |
||||
default: |
||||
return scopes, ErrInvalidQuotaTarget |
||||
} |
||||
} |
||||
@ -0,0 +1,28 @@ |
||||
package migrations |
||||
|
||||
import ( |
||||
. "github.com/grafana/grafana/pkg/services/sqlstore/migrator" |
||||
) |
||||
|
||||
func addQuotaMigration(mg *Migrator) { |
||||
|
||||
var quotaV1 = Table{ |
||||
Name: "quota", |
||||
Columns: []*Column{ |
||||
{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true}, |
||||
{Name: "org_id", Type: DB_BigInt, Nullable: true}, |
||||
{Name: "user_id", Type: DB_BigInt, Nullable: true}, |
||||
{Name: "target", Type: DB_NVarchar, Length: 255, Nullable: false}, |
||||
{Name: "limit", Type: DB_BigInt, Nullable: false}, |
||||
{Name: "created", Type: DB_DateTime, Nullable: false}, |
||||
{Name: "updated", Type: DB_DateTime, Nullable: false}, |
||||
}, |
||||
Indices: []*Index{ |
||||
{Cols: []string{"org_id", "user_id", "target"}, Type: UniqueIndex}, |
||||
}, |
||||
} |
||||
mg.AddMigration("create quota table v1", NewAddTableMigration(quotaV1)) |
||||
|
||||
//------- indexes ------------------
|
||||
addTableIndicesMigrations(mg, "v1", quotaV1) |
||||
} |
||||
@ -0,0 +1,239 @@ |
||||
package sqlstore |
||||
|
||||
import ( |
||||
"fmt" |
||||
"github.com/grafana/grafana/pkg/bus" |
||||
m "github.com/grafana/grafana/pkg/models" |
||||
"github.com/grafana/grafana/pkg/setting" |
||||
) |
||||
|
||||
func init() { |
||||
bus.AddHandler("sql", GetOrgQuotaByTarget) |
||||
bus.AddHandler("sql", GetOrgQuotas) |
||||
bus.AddHandler("sql", UpdateOrgQuota) |
||||
bus.AddHandler("sql", GetUserQuotaByTarget) |
||||
bus.AddHandler("sql", GetUserQuotas) |
||||
bus.AddHandler("sql", UpdateUserQuota) |
||||
bus.AddHandler("sql", GetGlobalQuotaByTarget) |
||||
} |
||||
|
||||
type targetCount struct { |
||||
Count int64 |
||||
} |
||||
|
||||
func GetOrgQuotaByTarget(query *m.GetOrgQuotaByTargetQuery) error { |
||||
quota := m.Quota{ |
||||
Target: query.Target, |
||||
OrgId: query.OrgId, |
||||
} |
||||
has, err := x.Get("a) |
||||
if err != nil { |
||||
return err |
||||
} else if has == false { |
||||
quota.Limit = query.Default |
||||
} |
||||
|
||||
//get quota used.
|
||||
rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(query.Target)) |
||||
resp := make([]*targetCount, 0) |
||||
if err := x.Sql(rawSql, query.OrgId).Find(&resp); err != nil { |
||||
return err |
||||
} |
||||
|
||||
query.Result = &m.OrgQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: quota.Limit, |
||||
OrgId: query.OrgId, |
||||
Used: resp[0].Count, |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func GetOrgQuotas(query *m.GetOrgQuotasQuery) error { |
||||
quotas := make([]*m.Quota, 0) |
||||
sess := x.Table("quota") |
||||
if err := sess.Where("org_id=? AND user_id=0", query.OrgId).Find("as); err != nil { |
||||
return err |
||||
} |
||||
|
||||
defaultQuotas := setting.Quota.Org.ToMap() |
||||
|
||||
seenTargets := make(map[string]bool) |
||||
for _, q := range quotas { |
||||
seenTargets[q.Target] = true |
||||
} |
||||
|
||||
for t, v := range defaultQuotas { |
||||
if _, ok := seenTargets[t]; !ok { |
||||
quotas = append(quotas, &m.Quota{ |
||||
OrgId: query.OrgId, |
||||
Target: t, |
||||
Limit: v, |
||||
}) |
||||
} |
||||
} |
||||
|
||||
result := make([]*m.OrgQuotaDTO, len(quotas)) |
||||
for i, q := range quotas { |
||||
//get quota used.
|
||||
rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target)) |
||||
resp := make([]*targetCount, 0) |
||||
if err := x.Sql(rawSql, q.OrgId).Find(&resp); err != nil { |
||||
return err |
||||
} |
||||
result[i] = &m.OrgQuotaDTO{ |
||||
Target: q.Target, |
||||
Limit: q.Limit, |
||||
OrgId: q.OrgId, |
||||
Used: resp[0].Count, |
||||
} |
||||
} |
||||
query.Result = result |
||||
return nil |
||||
} |
||||
|
||||
func UpdateOrgQuota(cmd *m.UpdateOrgQuotaCmd) error { |
||||
return inTransaction2(func(sess *session) error { |
||||
//Check if quota is already defined in the DB
|
||||
quota := m.Quota{ |
||||
Target: cmd.Target, |
||||
OrgId: cmd.OrgId, |
||||
} |
||||
has, err := sess.Get("a) |
||||
if err != nil { |
||||
return err |
||||
} |
||||
quota.Limit = cmd.Limit |
||||
if has == false { |
||||
//No quota in the DB for this target, so create a new one.
|
||||
if _, err := sess.Insert("a); err != nil { |
||||
return err |
||||
} |
||||
} else { |
||||
//update existing quota entry in the DB.
|
||||
if _, err := sess.Id(quota.Id).Update("a); err != nil { |
||||
return err |
||||
} |
||||
} |
||||
|
||||
return nil |
||||
}) |
||||
} |
||||
|
||||
func GetUserQuotaByTarget(query *m.GetUserQuotaByTargetQuery) error { |
||||
quota := m.Quota{ |
||||
Target: query.Target, |
||||
UserId: query.UserId, |
||||
} |
||||
has, err := x.Get("a) |
||||
if err != nil { |
||||
return err |
||||
} else if has == false { |
||||
quota.Limit = query.Default |
||||
} |
||||
|
||||
//get quota used.
|
||||
rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(query.Target)) |
||||
resp := make([]*targetCount, 0) |
||||
if err := x.Sql(rawSql, query.UserId).Find(&resp); err != nil { |
||||
return err |
||||
} |
||||
|
||||
query.Result = &m.UserQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: quota.Limit, |
||||
UserId: query.UserId, |
||||
Used: resp[0].Count, |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
|
||||
func GetUserQuotas(query *m.GetUserQuotasQuery) error { |
||||
quotas := make([]*m.Quota, 0) |
||||
sess := x.Table("quota") |
||||
if err := sess.Where("user_id=? AND org_id=0", query.UserId).Find("as); err != nil { |
||||
return err |
||||
} |
||||
|
||||
defaultQuotas := setting.Quota.User.ToMap() |
||||
|
||||
seenTargets := make(map[string]bool) |
||||
for _, q := range quotas { |
||||
seenTargets[q.Target] = true |
||||
} |
||||
|
||||
for t, v := range defaultQuotas { |
||||
if _, ok := seenTargets[t]; !ok { |
||||
quotas = append(quotas, &m.Quota{ |
||||
UserId: query.UserId, |
||||
Target: t, |
||||
Limit: v, |
||||
}) |
||||
} |
||||
} |
||||
|
||||
result := make([]*m.UserQuotaDTO, len(quotas)) |
||||
for i, q := range quotas { |
||||
//get quota used.
|
||||
rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where user_id=?", dialect.Quote(q.Target)) |
||||
resp := make([]*targetCount, 0) |
||||
if err := x.Sql(rawSql, q.UserId).Find(&resp); err != nil { |
||||
return err |
||||
} |
||||
result[i] = &m.UserQuotaDTO{ |
||||
Target: q.Target, |
||||
Limit: q.Limit, |
||||
UserId: q.UserId, |
||||
Used: resp[0].Count, |
||||
} |
||||
} |
||||
query.Result = result |
||||
return nil |
||||
} |
||||
|
||||
func UpdateUserQuota(cmd *m.UpdateUserQuotaCmd) error { |
||||
return inTransaction2(func(sess *session) error { |
||||
//Check if quota is already defined in the DB
|
||||
quota := m.Quota{ |
||||
Target: cmd.Target, |
||||
UserId: cmd.UserId, |
||||
} |
||||
has, err := sess.Get("a) |
||||
if err != nil { |
||||
return err |
||||
} |
||||
quota.Limit = cmd.Limit |
||||
if has == false { |
||||
//No quota in the DB for this target, so create a new one.
|
||||
if _, err := sess.Insert("a); err != nil { |
||||
return err |
||||
} |
||||
} else { |
||||
//update existing quota entry in the DB.
|
||||
if _, err := sess.Id(quota.Id).Update("a); err != nil { |
||||
return err |
||||
} |
||||
} |
||||
|
||||
return nil |
||||
}) |
||||
} |
||||
|
||||
func GetGlobalQuotaByTarget(query *m.GetGlobalQuotaByTargetQuery) error { |
||||
//get quota used.
|
||||
rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s", dialect.Quote(query.Target)) |
||||
resp := make([]*targetCount, 0) |
||||
if err := x.Sql(rawSql).Find(&resp); err != nil { |
||||
return err |
||||
} |
||||
|
||||
query.Result = &m.GlobalQuotaDTO{ |
||||
Target: query.Target, |
||||
Limit: query.Default, |
||||
Used: resp[0].Count, |
||||
} |
||||
|
||||
return nil |
||||
} |
||||
@ -0,0 +1,171 @@ |
||||
package sqlstore |
||||
|
||||
import ( |
||||
"testing" |
||||
|
||||
m "github.com/grafana/grafana/pkg/models" |
||||
"github.com/grafana/grafana/pkg/setting" |
||||
. "github.com/smartystreets/goconvey/convey" |
||||
) |
||||
|
||||
func TestQuotaCommandsAndQueries(t *testing.T) { |
||||
|
||||
Convey("Testing Qutoa commands & queries", t, func() { |
||||
InitTestDB(t) |
||||
userId := int64(1) |
||||
orgId := int64(0) |
||||
|
||||
setting.Quota = setting.QuotaSettings{ |
||||
Enabled: true, |
||||
Org: &setting.OrgQuota{ |
||||
User: 5, |
||||
Dashboard: 5, |
||||
DataSource: 5, |
||||
ApiKey: 5, |
||||
}, |
||||
User: &setting.UserQuota{ |
||||
Org: 5, |
||||
}, |
||||
Global: &setting.GlobalQuota{ |
||||
Org: 5, |
||||
User: 5, |
||||
Dashboard: 5, |
||||
DataSource: 5, |
||||
ApiKey: 5, |
||||
Session: 5, |
||||
}, |
||||
} |
||||
|
||||
// create a new org and add user_id 1 as admin.
|
||||
// we will then have an org with 1 user. and a user
|
||||
// with 1 org.
|
||||
userCmd := m.CreateOrgCommand{ |
||||
Name: "TestOrg", |
||||
UserId: 1, |
||||
} |
||||
err := CreateOrg(&userCmd) |
||||
So(err, ShouldBeNil) |
||||
orgId = userCmd.Result.Id |
||||
|
||||
Convey("Given saved org quota for users", func() { |
||||
orgCmd := m.UpdateOrgQuotaCmd{ |
||||
OrgId: orgId, |
||||
Target: "org_user", |
||||
Limit: 10, |
||||
} |
||||
err := UpdateOrgQuota(&orgCmd) |
||||
So(err, ShouldBeNil) |
||||
|
||||
Convey("Should be able to get saved quota by org id and target", func() { |
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 1} |
||||
err = GetOrgQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Limit, ShouldEqual, 10) |
||||
}) |
||||
Convey("Should be able to get default quota by org id and target", func() { |
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: 123, Target: "org_user", Default: 11} |
||||
err = GetOrgQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Limit, ShouldEqual, 11) |
||||
}) |
||||
Convey("Should be able to get used org quota when rows exist", func() { |
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: orgId, Target: "org_user", Default: 11} |
||||
err = GetOrgQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Used, ShouldEqual, 1) |
||||
}) |
||||
Convey("Should be able to get used org quota when no rows exist", func() { |
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: 2, Target: "org_user", Default: 11} |
||||
err = GetOrgQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Used, ShouldEqual, 0) |
||||
}) |
||||
Convey("Should be able to quota list for org", func() { |
||||
query := m.GetOrgQuotasQuery{OrgId: orgId} |
||||
err = GetOrgQuotas(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(len(query.Result), ShouldEqual, 4) |
||||
for _, res := range query.Result { |
||||
limit := 5 //default quota limit
|
||||
used := 0 |
||||
if res.Target == "org_user" { |
||||
limit = 10 //customized quota limit.
|
||||
used = 1 |
||||
} |
||||
So(res.Limit, ShouldEqual, limit) |
||||
So(res.Used, ShouldEqual, used) |
||||
|
||||
} |
||||
}) |
||||
}) |
||||
Convey("Given saved user quota for org", func() { |
||||
userQoutaCmd := m.UpdateUserQuotaCmd{ |
||||
UserId: userId, |
||||
Target: "org_user", |
||||
Limit: 10, |
||||
} |
||||
err := UpdateUserQuota(&userQoutaCmd) |
||||
So(err, ShouldBeNil) |
||||
|
||||
Convey("Should be able to get saved quota by user id and target", func() { |
||||
query := m.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 1} |
||||
err = GetUserQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Limit, ShouldEqual, 10) |
||||
}) |
||||
Convey("Should be able to get default quota by user id and target", func() { |
||||
query := m.GetUserQuotaByTargetQuery{UserId: 9, Target: "org_user", Default: 11} |
||||
err = GetUserQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Limit, ShouldEqual, 11) |
||||
}) |
||||
Convey("Should be able to get used user quota when rows exist", func() { |
||||
query := m.GetUserQuotaByTargetQuery{UserId: userId, Target: "org_user", Default: 11} |
||||
err = GetUserQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Used, ShouldEqual, 1) |
||||
}) |
||||
Convey("Should be able to get used user quota when no rows exist", func() { |
||||
query := m.GetUserQuotaByTargetQuery{UserId: 2, Target: "org_user", Default: 11} |
||||
err = GetUserQuotaByTarget(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(query.Result.Used, ShouldEqual, 0) |
||||
}) |
||||
Convey("Should be able to quota list for user", func() { |
||||
query := m.GetUserQuotasQuery{UserId: userId} |
||||
err = GetUserQuotas(&query) |
||||
|
||||
So(err, ShouldBeNil) |
||||
So(len(query.Result), ShouldEqual, 1) |
||||
So(query.Result[0].Limit, ShouldEqual, 10) |
||||
So(query.Result[0].Used, ShouldEqual, 1) |
||||
}) |
||||
}) |
||||
|
||||
Convey("Should be able to global user quota", func() { |
||||
query := m.GetGlobalQuotaByTargetQuery{Target: "user", Default: 5} |
||||
err = GetGlobalQuotaByTarget(&query) |
||||
So(err, ShouldBeNil) |
||||
|
||||
So(query.Result.Limit, ShouldEqual, 5) |
||||
So(query.Result.Used, ShouldEqual, 0) |
||||
}) |
||||
Convey("Should be able to global org quota", func() { |
||||
query := m.GetGlobalQuotaByTargetQuery{Target: "org", Default: 5} |
||||
err = GetGlobalQuotaByTarget(&query) |
||||
So(err, ShouldBeNil) |
||||
|
||||
So(query.Result.Limit, ShouldEqual, 5) |
||||
So(query.Result.Used, ShouldEqual, 1) |
||||
}) |
||||
}) |
||||
} |
||||
@ -0,0 +1,94 @@ |
||||
package setting |
||||
|
||||
import ( |
||||
"reflect" |
||||
) |
||||
|
||||
type OrgQuota struct { |
||||
User int64 `target:"org_user"` |
||||
DataSource int64 `target:"data_source"` |
||||
Dashboard int64 `target:"dashboard"` |
||||
ApiKey int64 `target:"api_key"` |
||||
} |
||||
|
||||
type UserQuota struct { |
||||
Org int64 `target:"org_user"` |
||||
} |
||||
|
||||
type GlobalQuota struct { |
||||
Org int64 `target:"org"` |
||||
User int64 `target:"user"` |
||||
DataSource int64 `target:"data_source"` |
||||
Dashboard int64 `target:"dashboard"` |
||||
ApiKey int64 `target:"api_key"` |
||||
Session int64 `target:"-"` |
||||
} |
||||
|
||||
func (q *OrgQuota) ToMap() map[string]int64 { |
||||
return quotaToMap(*q) |
||||
} |
||||
|
||||
func (q *UserQuota) ToMap() map[string]int64 { |
||||
return quotaToMap(*q) |
||||
} |
||||
|
||||
func (q *GlobalQuota) ToMap() map[string]int64 { |
||||
return quotaToMap(*q) |
||||
} |
||||
|
||||
func quotaToMap(q interface{}) map[string]int64 { |
||||
qMap := make(map[string]int64) |
||||
typ := reflect.TypeOf(q) |
||||
val := reflect.ValueOf(q) |
||||
|
||||
for i := 0; i < typ.NumField(); i++ { |
||||
field := typ.Field(i) |
||||
name := field.Tag.Get("target") |
||||
if name == "" { |
||||
name = field.Name |
||||
} |
||||
if name == "-" { |
||||
continue |
||||
} |
||||
value := val.Field(i) |
||||
qMap[name] = value.Int() |
||||
} |
||||
return qMap |
||||
} |
||||
|
||||
type QuotaSettings struct { |
||||
Enabled bool |
||||
Org *OrgQuota |
||||
User *UserQuota |
||||
Global *GlobalQuota |
||||
} |
||||
|
||||
func readQuotaSettings() { |
||||
// set global defaults.
|
||||
quota := Cfg.Section("quota") |
||||
Quota.Enabled = quota.Key("enabled").MustBool(false) |
||||
|
||||
// per ORG Limits
|
||||
Quota.Org = &OrgQuota{ |
||||
User: quota.Key("org_user").MustInt64(10), |
||||
DataSource: quota.Key("org_data_source").MustInt64(10), |
||||
Dashboard: quota.Key("org_dashboard").MustInt64(10), |
||||
ApiKey: quota.Key("org_api_key").MustInt64(10), |
||||
} |
||||
|
||||
// per User limits
|
||||
Quota.User = &UserQuota{ |
||||
Org: quota.Key("user_org").MustInt64(10), |
||||
} |
||||
|
||||
// Global Limits
|
||||
Quota.Global = &GlobalQuota{ |
||||
User: quota.Key("global_user").MustInt64(-1), |
||||
Org: quota.Key("global_org").MustInt64(-1), |
||||
DataSource: quota.Key("global_data_source").MustInt64(-1), |
||||
Dashboard: quota.Key("global_dashboard").MustInt64(-1), |
||||
ApiKey: quota.Key("global_api_key").MustInt64(-1), |
||||
Session: quota.Key("global_session").MustInt64(-1), |
||||
} |
||||
|
||||
} |
||||
Loading…
Reference in new issue