mirror of https://github.com/grafana/grafana
commit
3912eb7b26
@ -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 := setting.Quota.Org.ToMap()[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 := setting.Quota.User.ToMap()[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