IntenseWebs/grafana
3
0
Fork 0
mirror of https://github.com/grafana/grafana.git synced 2025-02-25 18:55:37 -06:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'quotas' of https://github.com/raintank/grafana into raintank-quotas

Browse Source
This commit is contained in:
Torkel Ödegaard 2015-09-15 14:23:13 +02:00
commit 6a30511fc4
14 changed files with 1051 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
conf/defaults.ini
View File

@ -87,6 +87,7 @@ cookie_secure = false
# Session life time, default is 86400
session_life_time = 86400
gc_interval_time = 86400
#################################### Analytics ####################################
[analytics]
@ -253,4 +254,37 @@ exchange = grafana_events
enabled = false
path = /var/lib/grafana/dashboards
#################################### Usage Quotas ##########################
[quota]
enabled = false
#### set quotas to -1 to make unlimited. ####
# limit number of users per Org.
org_user = 10
# limit number of dashboards per Org.
org_dashboard = 100
# limit number of data_sources per Org.
org_data_source = 10
# limit number of api_keys per Org.
org_api_key = 10
# limit number of orgs a user can create.
user_org = 10
# Global limit of users.
global_user = -1
# global limit of orgs.
global_org = -1
# global limit of dashboards
global_dashboard = -1
# global limit of api_keys
global_api_key = -1
# global limit on number of logged in users.
global_session = -1

32
pkg/api/api.go
View File

@ -14,13 +14,14 @@ func Register(r *macaron.Macaron) {
reqGrafanaAdmin := middleware.Auth(&middleware.AuthOptions{ReqSignedIn: true, ReqGrafanaAdmin: true})
reqEditorRole := middleware.RoleAuth(m.ROLE_EDITOR, m.ROLE_ADMIN)
regOrgAdmin := middleware.RoleAuth(m.ROLE_ADMIN)
quota := middleware.Quota
bind := binding.Bind
// not logged in views
r.Get("/", reqSignedIn, Index)
r.Get("/logout", Logout)
r.Post("/login", bind(dtos.LoginCommand{}), wrap(LoginPost))
r.Get("/login/:name", OAuthLogin)
r.Post("/login", quota("session"), bind(dtos.LoginCommand{}), wrap(LoginPost))
r.Get("/login/:name", quota("session"), OAuthLogin)
r.Get("/login", LoginView)
r.Get("/invite/:code", Index)
@ -44,7 +45,7 @@ func Register(r *macaron.Macaron) {
// sign up
r.Get("/signup", Index)
r.Get("/api/user/signup/options", wrap(GetSignUpOptions))
r.Post("/api/user/signup", bind(dtos.SignUpForm{}), wrap(SignUp))
r.Post("/api/user/signup", quota("user"), bind(dtos.SignUpForm{}), wrap(SignUp))
r.Post("/api/user/signup/step2", bind(dtos.SignUpStep2Form{}), wrap(SignUpStep2))
// invited
@ -66,7 +67,7 @@ func Register(r *macaron.Macaron) {
r.Get("/api/snapshots-delete/:key", DeleteDashboardSnapshot)
// api renew session based on remember cookie
r.Get("/api/login/ping", LoginApiPing)
r.Get("/api/login/ping", quota("session"), LoginApiPing)
// authed api
r.Group("/api", func() {
@ -80,6 +81,7 @@ func Register(r *macaron.Macaron) {
r.Post("/stars/dashboard/:id", wrap(StarDashboard))
r.Delete("/stars/dashboard/:id", wrap(UnstarDashboard))
r.Put("/password", bind(m.ChangeUserPasswordCommand{}), wrap(ChangeUserPassword))
r.Get("/quotas", wrap(GetUserQuotas))
})
// users (admin permission required)
@ -90,24 +92,29 @@ func Register(r *macaron.Macaron) {
r.Put("/:id", bind(m.UpdateUserCommand{}), wrap(UpdateUser))
}, reqGrafanaAdmin)
// current org
// org information available to all users.
r.Group("/org", func() {
r.Get("/", wrap(GetOrgCurrent))
r.Get("/quotas", wrap(GetOrgQuotas))
})
// current org
r.Group("/org", func() {
r.Put("/", bind(dtos.UpdateOrgForm{}), wrap(UpdateOrgCurrent))
r.Put("/address", bind(dtos.UpdateOrgAddressForm{}), wrap(UpdateOrgAddressCurrent))
r.Post("/users", bind(m.AddOrgUserCommand{}), wrap(AddOrgUserToCurrentOrg))
r.Post("/users", quota("user"), bind(m.AddOrgUserCommand{}), wrap(AddOrgUserToCurrentOrg))
r.Get("/users", wrap(GetOrgUsersForCurrentOrg))
r.Patch("/users/:userId", bind(m.UpdateOrgUserCommand{}), wrap(UpdateOrgUserForCurrentOrg))
r.Delete("/users/:userId", wrap(RemoveOrgUserForCurrentOrg))
// invites
r.Get("/invites", wrap(GetPendingOrgInvites))
r.Post("/invites", bind(dtos.AddInviteForm{}), wrap(AddOrgInvite))
r.Post("/invites", quota("user"), bind(dtos.AddInviteForm{}), wrap(AddOrgInvite))
r.Patch("/invites/:code/revoke", wrap(RevokeInvite))
}, regOrgAdmin)
// create new org
r.Post("/orgs", bind(m.CreateOrgCommand{}), wrap(CreateOrg))
r.Post("/orgs", quota("org"), bind(m.CreateOrgCommand{}), wrap(CreateOrg))
// search all orgs
r.Get("/orgs", reqGrafanaAdmin, wrap(SearchOrgs))
@ -122,19 +129,21 @@ func Register(r *macaron.Macaron) {
r.Post("/users", bind(m.AddOrgUserCommand{}), wrap(AddOrgUser))
r.Patch("/users/:userId", bind(m.UpdateOrgUserCommand{}), wrap(UpdateOrgUser))
r.Delete("/users/:userId", wrap(RemoveOrgUser))
r.Get("/quotas", wrap(GetOrgQuotas))
r.Put("/quotas/:target", bind(m.UpdateOrgQuotaCmd{}), wrap(UpdateOrgQuota))
}, reqGrafanaAdmin)
// auth api keys
r.Group("/auth/keys", func() {
r.Get("/", wrap(GetApiKeys))
r.Post("/", bind(m.AddApiKeyCommand{}), wrap(AddApiKey))
r.Post("/", quota("api_key"), bind(m.AddApiKeyCommand{}), wrap(AddApiKey))
r.Delete("/:id", wrap(DeleteApiKey))
}, regOrgAdmin)
// Data sources
r.Group("/datasources", func() {
r.Get("/", GetDataSources)
r.Post("/", bind(m.AddDataSourceCommand{}), AddDataSource)
r.Post("/", quota("data_source"), bind(m.AddDataSourceCommand{}), AddDataSource)
r.Put("/:id", bind(m.UpdateDataSourceCommand{}), UpdateDataSource)
r.Delete("/:id", DeleteDataSource)
r.Get("/:id", GetDataSourceById)
@ -159,6 +168,7 @@ func Register(r *macaron.Macaron) {
// metrics
r.Get("/metrics/test", GetTestMetrics)
}, reqSignedIn)
// admin api
@ -168,6 +178,8 @@ func Register(r *macaron.Macaron) {
r.Put("/users/:id/password", bind(dtos.AdminUpdateUserPasswordForm{}), AdminUpdateUserPassword)
r.Put("/users/:id/permissions", bind(dtos.AdminUpdateUserPermissionsForm{}), AdminUpdateUserPermissions)
r.Delete("/users/:id", AdminDeleteUser)
r.Get("/users/:id/quotas", wrap(GetUserQuotas))
r.Put("/users/:id/quotas/:target", bind(m.UpdateUserQuotaCmd{}), wrap(UpdateUserQuota))
}, reqGrafanaAdmin)
// rendering

13
pkg/api/dashboard.go
View File

@ -86,6 +86,19 @@ func DeleteDashboard(c *middleware.Context) {
func PostDashboard(c *middleware.Context, cmd m.SaveDashboardCommand) {
cmd.OrgId = c.OrgId
dash := cmd.GetDashboardModel()
if dash.Id == 0 {
limitReached, err := middleware.QuotaReached(c, "dashboard")
if err != nil {
c.JsonApiErr(500, "failed to get quota", err)
return
}
if limitReached {
c.JsonApiErr(403, "Quota reached", nil)
return
}
}
err := bus.Dispatch(&cmd)
if err != nil {
if err == m.ErrDashboardWithSameNameExists {

10
pkg/api/login_oauth.go
View File

@ -74,7 +74,15 @@ func OAuthLogin(ctx *middleware.Context) {
ctx.Redirect(setting.AppSubUrl + "/login")
return
}
limitReached, err := middleware.QuotaReached(ctx, "user")
if err != nil {
ctx.Handle(500, "Failed to get user quota", err)
return
}
if limitReached {
ctx.Redirect(setting.AppSubUrl + "/login")
return
}
cmd := m.CreateUserCommand{
Login: userInfo.Email,
Email: userInfo.Email,

68
pkg/api/quota.go Normal file
View File

@ -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")
}

94
pkg/middleware/middleware.go
View File

@ -1,6 +1,7 @@
package middleware
import (
"fmt"
"strconv"
"strings"
@ -253,3 +254,96 @@ func (ctx *Context) JsonApiErr(status int, message string, err error) {
ctx.JSON(status, resp)
}
func Quota(target string) macaron.Handler {
return func(c *Context) {
limitReached, err := QuotaReached(c, target)
if err != nil {
c.JsonApiErr(500, "failed to get quota", err)
return
}
if limitReached {
c.JsonApiErr(403, fmt.Sprintf("%s Quota reached", target), nil)
return
}
}
}
func QuotaReached(c *Context, target string) (bool, error) {
if !setting.Quota.Enabled {
return false, nil
}
// get the list of scopes that this target is valid for. Org, User, Global
scopes, err := m.GetQuotaScopes(target)
if err != nil {
return false, err
}
log.Info(fmt.Sprintf("checking quota for %s in scopes %v", target, scopes))
for _, scope := range scopes {
log.Info(fmt.Sprintf("checking scope %s", scope.Name))
switch scope.Name {
case "global":
if scope.DefaultLimit < 0 {
continue
}
if scope.DefaultLimit == 0 {
return true, nil
}
if target == "session" {
usedSessions := sessionManager.Count()
if int64(usedSessions) > scope.DefaultLimit {
log.Info(fmt.Sprintf("%d sessions active, limit is %d", usedSessions, scope.DefaultLimit))
return true, nil
}
continue
}
query := m.GetGlobalQuotaByTargetQuery{Target: scope.Target}
if err := bus.Dispatch(&query); err != nil {
return true, err
}
if query.Result.Used >= scope.DefaultLimit {
return true, nil
}
case "org":
if !c.IsSignedIn {
continue
}
query := m.GetOrgQuotaByTargetQuery{OrgId: c.OrgId, Target: scope.Target, Default: scope.DefaultLimit}
if err := bus.Dispatch(&query); err != nil {
return true, err
}
if query.Result.Limit < 0 {
continue
}
if query.Result.Limit == 0 {
return true, nil
}
if query.Result.Used >= query.Result.Limit {
return true, nil
}
case "user":
if !c.IsSignedIn || c.UserId == 0 {
continue
}
query := m.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit}
if err := bus.Dispatch(&query); err != nil {
return true, err
}
if query.Result.Limit < 0 {
continue
}
if query.Result.Limit == 0 {
return true, nil
}
if query.Result.Used >= query.Result.Limit {
return true, nil
}
}
}
return false, nil
}

144
pkg/middleware/quota_test.go Normal file
View File

@ -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)
})
})
})
}

130
pkg/models/quotas.go Normal file
View File

@ -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
}
}

1
pkg/services/sqlstore/migrations/migrations.go
View File

@ -17,6 +17,7 @@ func AddMigrations(mg *Migrator) {
addDataSourceMigration(mg)
addApiKeyMigrations(mg)
addDashboardSnapshotMigrations(mg)
addQuotaMigration(mg)
}
func addMigrationLogMigrations(mg *Migrator) {

28
pkg/services/sqlstore/migrations/quota_mig.go Normal file
View File

@ -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)
}

239
pkg/services/sqlstore/quota.go Normal file
View File

@ -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(&quota)
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(&quotas); 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(&quota)
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(&quota); err != nil {
return err
}
} else {
//update existing quota entry in the DB.
if _, err := sess.Id(quota.Id).Update(&quota); 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(&quota)
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(&quotas); 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(&quota)
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(&quota); err != nil {
return err
}
} else {
//update existing quota entry in the DB.
if _, err := sess.Id(quota.Id).Update(&quota); 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
}

171
pkg/services/sqlstore/quota_test.go Normal file
View File

@ -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)
})
})
}

4
pkg/setting/setting.go
View File

@ -127,6 +127,9 @@ var (
// SMTP email settings
Smtp SmtpSettings
// QUOTA
Quota QuotaSettings
)
type CommandLineArgs struct {
@ -458,6 +461,7 @@ func NewConfigContext(args *CommandLineArgs) error {
readSessionConfig()
readSmtpSettings()
readQuotaSettings()
if VerifyEmailEnabled && !Smtp.Enabled {
log.Warn("require_email_validation is enabled but smpt is disabled")

94
pkg/setting/setting_quota.go Normal file
View File

@ -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),
}
}