diff --git a/conf/defaults.ini b/conf/defaults.ini index 941813fb4c6..d640820fe65 100644 --- a/conf/defaults.ini +++ b/conf/defaults.ini @@ -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 diff --git a/pkg/api/api.go b/pkg/api/api.go index 6c56e6aa9f1..27eb3c749db 100644 --- a/pkg/api/api.go +++ b/pkg/api/api.go @@ -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 diff --git a/pkg/api/dashboard.go b/pkg/api/dashboard.go index a10c3c92f96..d7f9f2cd740 100644 --- a/pkg/api/dashboard.go +++ b/pkg/api/dashboard.go @@ -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 { diff --git a/pkg/api/login_oauth.go b/pkg/api/login_oauth.go index 796599df864..4244feef664 100644 --- a/pkg/api/login_oauth.go +++ b/pkg/api/login_oauth.go @@ -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, diff --git a/pkg/api/quota.go b/pkg/api/quota.go new file mode 100644 index 00000000000..490a8e90729 --- /dev/null +++ b/pkg/api/quota.go @@ -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") +} diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go index 8704ec5a787..8a61bdb3a0b 100644 --- a/pkg/middleware/middleware.go +++ b/pkg/middleware/middleware.go @@ -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 +} diff --git a/pkg/middleware/quota_test.go b/pkg/middleware/quota_test.go new file mode 100644 index 00000000000..ae6b414b248 --- /dev/null +++ b/pkg/middleware/quota_test.go @@ -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) + }) + + }) + + }) +} diff --git a/pkg/models/quotas.go b/pkg/models/quotas.go new file mode 100644 index 00000000000..85159d830d1 --- /dev/null +++ b/pkg/models/quotas.go @@ -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 + } +} diff --git a/pkg/services/sqlstore/migrations/migrations.go b/pkg/services/sqlstore/migrations/migrations.go index 70acc9ef736..8f7054d3959 100644 --- a/pkg/services/sqlstore/migrations/migrations.go +++ b/pkg/services/sqlstore/migrations/migrations.go @@ -17,6 +17,7 @@ func AddMigrations(mg *Migrator) { addDataSourceMigration(mg) addApiKeyMigrations(mg) addDashboardSnapshotMigrations(mg) + addQuotaMigration(mg) } func addMigrationLogMigrations(mg *Migrator) { diff --git a/pkg/services/sqlstore/migrations/quota_mig.go b/pkg/services/sqlstore/migrations/quota_mig.go new file mode 100644 index 00000000000..d877bb4c3c1 --- /dev/null +++ b/pkg/services/sqlstore/migrations/quota_mig.go @@ -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) +} diff --git a/pkg/services/sqlstore/quota.go b/pkg/services/sqlstore/quota.go new file mode 100644 index 00000000000..53ea8889c56 --- /dev/null +++ b/pkg/services/sqlstore/quota.go @@ -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 +} diff --git a/pkg/services/sqlstore/quota_test.go b/pkg/services/sqlstore/quota_test.go new file mode 100644 index 00000000000..5ef618e166d --- /dev/null +++ b/pkg/services/sqlstore/quota_test.go @@ -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) + }) + }) +} diff --git a/pkg/setting/setting.go b/pkg/setting/setting.go index b706d5a0957..8d62e475fc1 100644 --- a/pkg/setting/setting.go +++ b/pkg/setting/setting.go @@ -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") diff --git a/pkg/setting/setting_quota.go b/pkg/setting/setting_quota.go new file mode 100644 index 00000000000..49769d9930f --- /dev/null +++ b/pkg/setting/setting_quota.go @@ -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), + } + +}