From fa45742abcc5ca4c0cc11b487baa33ec0b0e9f6b Mon Sep 17 00:00:00 2001
From: Eric Leijonmarck <eric.leijonmarck@gmail.com>
Date: Tue, 18 Oct 2022 12:26:38 +0100
Subject: [PATCH] Quota(fix): remove service accounts from quota count (#56700)

---
 pkg/services/sqlstore/quota.go      | 26 +++++++++++++++++++++++---
 pkg/services/sqlstore/quota_test.go | 15 +++++++++++----
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/pkg/services/sqlstore/quota.go b/pkg/services/sqlstore/quota.go
index 6f7fc67f28b..a28dba881d7 100644
--- a/pkg/services/sqlstore/quota.go
+++ b/pkg/services/sqlstore/quota.go
@@ -41,7 +41,14 @@ func (ss *SQLStore) GetOrgQuotaByTarget(ctx context.Context, query *models.GetOr
 			if query.Target == dashboardTarget {
 				rawSQL += fmt.Sprintf(" AND is_folder=%s", dialect.BooleanStr(false))
 			}
-
+			// need to account for removing service accounts from the user table
+			if query.Target == "org_user" {
+				rawSQL = fmt.Sprintf("SELECT COUNT(*) as count from (select user_id from %s where org_id=? AND user_id IN (SELECT id as user_id FROM %s WHERE is_service_account=%s)) as subq",
+					dialect.Quote(query.Target),
+					dialect.Quote("user"),
+					dialect.BooleanStr(false),
+				)
+			}
 			resp := make([]*targetCount, 0)
 			if err := sess.SQL(rawSQL, query.OrgId).Find(&resp); err != nil {
 				return err
@@ -87,9 +94,19 @@ func (ss *SQLStore) GetOrgQuotas(ctx context.Context, query *models.GetOrgQuotas
 		result := make([]*models.OrgQuotaDTO, len(quotas))
 		for i, q := range quotas {
 			var used int64
+			var rawSQL string
 			if q.Target != alertRuleTarget || query.UnifiedAlertingEnabled {
 				// get quota used.
-				rawSQL := fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
+				rawSQL = fmt.Sprintf("SELECT COUNT(*) as count from %s where org_id=?", dialect.Quote(q.Target))
+
+				// need to account for removing service accounts from the user table
+				if q.Target == "org_user" {
+					rawSQL = fmt.Sprintf("SELECT COUNT(*) as count from (select user_id from %s where org_id=? AND user_id IN (SELECT id as user_id FROM %s WHERE is_service_account=%s)) as subq",
+						dialect.Quote(q.Target),
+						dialect.Quote("user"),
+						dialect.BooleanStr(false),
+					)
+				}
 				resp := make([]*targetCount, 0)
 				if err := sess.SQL(rawSQL, q.OrgId).Find(&resp); err != nil {
 					return err
@@ -276,7 +293,10 @@ func (ss *SQLStore) GetGlobalQuotaByTarget(ctx context.Context, query *models.Ge
 			if query.Target == dashboardTarget {
 				rawSQL += fmt.Sprintf(" WHERE is_folder=%s", dialect.BooleanStr(false))
 			}
-
+			// removing service accounts from count
+			if query.Target == dialect.Quote("user") {
+				rawSQL += fmt.Sprintf(" WHERE is_service_account=%s", dialect.BooleanStr(false))
+			}
 			resp := make([]*targetCount, 0)
 			if err := sess.SQL(rawSQL).Find(&resp); err != nil {
 				return err
diff --git a/pkg/services/sqlstore/quota_test.go b/pkg/services/sqlstore/quota_test.go
index 79b37903359..e58b42adf8d 100644
--- a/pkg/services/sqlstore/quota_test.go
+++ b/pkg/services/sqlstore/quota_test.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/stretchr/testify/require"
 )
@@ -40,16 +41,22 @@ func TestIntegrationQuotaCommandsAndQueries(t *testing.T) {
 			AlertRule:  5,
 		},
 	}
-
+	createUserCmd := user.CreateUserCommand{
+		Name:         "TestUser",
+		OrgID:        orgId,
+		SkipOrgSetup: true,
+	}
+	user, err := sqlStore.CreateUser(context.Background(), createUserCmd)
+	require.NoError(t, err)
 	// 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 := models.CreateOrgCommand{
 		Name:   "TestOrg",
-		UserId: 1,
+		UserId: user.ID,
 	}
 
-	err := sqlStore.CreateOrg(context.Background(), &userCmd)
+	err = sqlStore.CreateOrg(context.Background(), &userCmd)
 	require.NoError(t, err)
 	orgId = userCmd.Result.Id
 
@@ -198,7 +205,7 @@ func TestIntegrationQuotaCommandsAndQueries(t *testing.T) {
 		require.NoError(t, err)
 
 		require.Equal(t, int64(5), query.Result.Limit)
-		require.Equal(t, int64(0), query.Result.Used)
+		require.Equal(t, int64(1), query.Result.Used)
 	})
 
 	t.Run("Should be able to global org quota", func(t *testing.T) {