enhance quota support.

now includes: - perOrg (users, dashboards, datasources, api_keys) - perUser (orgs) - global (users, orgs, dashboards, datasources, api_keys, sessions)
2015-09-11 23:17:10 +08:00
parent 47bf1bd21a
commit 6488324cf1
10 changed files with 485 additions and 101 deletions
@@ -10,14 +10,15 @@ func addQuotaMigration(mg *Migrator) {
 		Name: "quota",
 		Columns: []*Column{
 			{Name: "id", Type: DB_BigInt, IsPrimaryKey: true, IsAutoIncrement: true},
-			{Name: "org_id", Type: DB_BigInt, Nullable: false},
+			{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", "target"}, Type: UniqueIndex},
+			{Cols: []string{"org_id", "user_id", "target"}, Type: UniqueIndex},
 		},
 	}
 	mg.AddMigration("create quota table v1", NewAddTableMigration(quotaV1))
@@ -8,16 +8,20 @@ import (
 )

 func init() {
-	bus.AddHandler("sql", GetQuotaByTarget)
-	bus.AddHandler("sql", GetQuotas)
-	bus.AddHandler("sql", UpdateQuota)
+	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 GetQuotaByTarget(query *m.GetQuotaByTargetQuery) error {
+func GetOrgQuotaByTarget(query *m.GetOrgQuotaByTargetQuery) error {
 	quota := m.Quota{
 		Target: query.Target,
 		OrgId:  query.OrgId,
@@ -26,17 +30,17 @@ func GetQuotaByTarget(query *m.GetQuotaByTargetQuery) error {
 	if err != nil {
 		return err
 	} else if has == false {
-		quota.Limit = setting.Quota.Default[string(query.Target)]
+		quota.Limit = query.Default
 	}

 	//get quota used.
-	rawSql := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(string(query.Target)))
+	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.QuotaDTO{
+	query.Result = &m.OrgQuotaDTO{
 		Target: query.Target,
 		Limit:  quota.Limit,
 		OrgId:  query.OrgId,
@@ -46,36 +50,39 @@ func GetQuotaByTarget(query *m.GetQuotaByTargetQuery) error {
 	return nil
 }

-func GetQuotas(query *m.GetQuotasQuery) error {
+func GetOrgQuotas(query *m.GetOrgQuotasQuery) error {
 	quotas := make([]*m.Quota, 0)
 	sess := x.Table("quota")
-	if err := sess.Where("org_id=?", query.OrgId).Find(&quotas); err != nil {
+	if err := sess.Where("org_id=? AND user_id=0", query.OrgId).Find(&quotas); err != nil {
 		return err
 	}

-	seenTargets := make(map[m.QuotaTarget]bool)
+	defaultQuotas := m.QuotaToMap(setting.Quota.Org)
+
+	seenTargets := make(map[string]bool)
 	for _, q := range quotas {
 		seenTargets[q.Target] = true
 	}

-	for t, v := range setting.Quota.Default {
-		if _, ok := seenTargets[m.QuotaTarget(t)]; !ok {
+	for t, v := range defaultQuotas {
+		if _, ok := seenTargets[t]; !ok {
 			quotas = append(quotas, &m.Quota{
 				OrgId:  query.OrgId,
-				Target: m.QuotaTarget(t),
+				Target: t,
 				Limit:  v,
 			})
 		}
 	}
-	result := make([]*m.QuotaDTO, len(quotas))
+
+	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(string(q.Target)))
+		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.QuotaDTO{
+		result[i] = &m.OrgQuotaDTO{
 			Target: q.Target,
 			Limit:  q.Limit,
 			OrgId:  q.OrgId,
@@ -86,7 +93,7 @@ func GetQuotas(query *m.GetQuotasQuery) error {
 	return nil
 }

-func UpdateQuota(cmd *m.UpdateQuotaCmd) error {
+func UpdateOrgQuota(cmd *m.UpdateOrgQuotaCmd) error {
 	return inTransaction2(func(sess *session) error {
 		//Check if quota is already defined in the DB
 		quota := m.Quota{
@@ -113,3 +120,120 @@ func UpdateQuota(cmd *m.UpdateQuotaCmd) error {
 		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 := m.QuotaToMap(setting.Quota.User)
+
+	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
+}