From 3d4d8225280dd2fe9292d98f86099281dec50866 Mon Sep 17 00:00:00 2001
From: woodsaj <anthony@monkey.id.au>
Date: Tue, 21 Jul 2015 18:30:31 +0800
Subject: [PATCH] implement updateQuota function

---
 pkg/api/quota.go               |  5 +++++
 pkg/models/quotas.go           | 11 +++++++++++
 pkg/services/sqlstore/quota.go | 26 +++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/pkg/api/quota.go b/pkg/api/quota.go
index 13e0e8c9781..1811270c3a4 100644
--- a/pkg/api/quota.go
+++ b/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)
 	}
diff --git a/pkg/models/quotas.go b/pkg/models/quotas.go
index d16bd0325cc..27ab74cb1e2 100644
--- a/pkg/models/quotas.go
+++ b/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
diff --git a/pkg/services/sqlstore/quota.go b/pkg/services/sqlstore/quota.go
index 244c49cc023..a5ce44fc4b6 100644
--- a/pkg/services/sqlstore/quota.go
+++ b/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
+	})
 }