Browse Source

implement updateQuota function

woodsaj 10 năm trước cách đây
mục cha
commit
3d4d822528
3 tập tin đã thay đổi với 41 bổ sung1 xóa
  1. 5 0
      pkg/api/quota.go
  2. 11 0
      pkg/models/quotas.go
  3. 25 1
      pkg/services/sqlstore/quota.go

+ 5 - 0
pkg/api/quota.go

@@ -19,6 +19,11 @@ func GetOrgQuotas(c *middleware.Context) Response {
 func UpdateOrgQuota(c *middleware.Context, cmd m.UpdateQuotaCmd) Response {
 	cmd.OrgId = c.ParamsInt64(":orgId")
 	cmd.Target = m.QuotaTarget(c.Params(":target"))
+
+	if !cmd.Target.IsValid() {
+		return ApiError(404, "Invalid quota target", nil)
+	}
+
 	if err := bus.Dispatch(&cmd); err != nil {
 		return ApiError(500, "Failed to update org quotas", err)
 	}

+ 11 - 0
pkg/models/quotas.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"errors"
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/setting"
 	"time"
@@ -16,6 +17,13 @@ const (
 	QUOTA_COLLECTOR  QuotaTarget = "collector"
 )
 
+var ErrInvalidQuotaTarget = errors.New("Invalid quota target")
+
+func (q QuotaTarget) IsValid() bool {
+	_, ok := DefaultQuotas[q]
+	return ok
+}
+
 // defaults are set from settings package.
 var DefaultQuotas map[QuotaTarget]int64
 
@@ -64,6 +72,9 @@ type UpdateQuotaCmd struct {
 }
 
 func QuotaReached(org_id int64, target QuotaTarget) (bool, error) {
+	if !target.IsValid() {
+		return true, ErrInvalidQuotaTarget
+	}
 	query := GetQuotaByTargetQuery{OrgId: org_id, Target: target}
 	if err := bus.Dispatch(&query); err != nil {
 		return true, err

+ 25 - 1
pkg/services/sqlstore/quota.go

@@ -86,5 +86,29 @@ func GetQuotas(query *m.GetQuotasQuery) error {
 }
 
 func UpdateQuota(cmd *m.UpdateQuotaCmd) error {
-	return nil
+	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
+	})
 }