QuotaService: refactor to use ReplDB for Get queries (#91333)

* Feature (quota service): Use ReplDB for quota service Gets

This adds the replDB to the quota service, as well as some more test helper functions to simplify updating tests. My intent is that the helper functions can be removed when this is fully rolled out (or not) and we're consistently using the ReplDB interface (or not!)

* test updates

(cherry picked from commit 299c142f6a)
This commit is contained in:
Kristin Laemmert
2024-08-21 15:23:47 +01:00
committed by Andreas Christou
parent 32c95a1a70
commit f4b646da29
20 changed files with 74 additions and 53 deletions
+1 -1
View File
@@ -906,7 +906,7 @@ func TestIntegration_SQLStore_RemoveOrgUser(t *testing.T) {
func createOrgAndUserSvc(t *testing.T, store db.DB, cfg *setting.Cfg) (org.Service, user.Service) {
t.Helper()
quotaService := quotaimpl.ProvideService(store, cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), cfg)
orgService, err := ProvideService(store, cfg, quotaService)
require.NoError(t, err)
usrSvc, err := userimpl.ProvideService(
+1 -1
View File
@@ -53,7 +53,7 @@ type service struct {
targetToSrv *quota.TargetToSrv
}
func ProvideService(db db.DB, cfg *setting.Cfg) quota.Service {
func ProvideService(db db.ReplDB, cfg *setting.Cfg) quota.Service {
logger := log.New("quota_service")
s := service{
store: &sqlStore{db: db, logger: logger},
+5 -5
View File
@@ -16,12 +16,12 @@ type store interface {
}
type sqlStore struct {
db db.DB
db db.ReplDB
logger log.Logger
}
func (ss *sqlStore) DeleteByUser(ctx quota.Context, userID int64) error {
return ss.db.WithDbSession(ctx, func(sess *db.Session) error {
return ss.db.DB().WithDbSession(ctx, func(sess *db.Session) error {
var rawSQL = "DELETE FROM quota WHERE user_id = ?"
_, err := sess.Exec(rawSQL, userID)
return err
@@ -54,7 +54,7 @@ func (ss *sqlStore) Get(ctx quota.Context, scopeParams *quota.ScopeParameters) (
}
func (ss *sqlStore) Update(ctx quota.Context, cmd *quota.UpdateQuotaCmd) error {
return ss.db.WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
return ss.db.DB().WithTransactionalDbSession(ctx, func(sess *sqlstore.DBSession) error {
// Check if quota is already defined in the DB
quota := quota.Quota{
Target: cmd.Target,
@@ -87,7 +87,7 @@ func (ss *sqlStore) Update(ctx quota.Context, cmd *quota.UpdateQuotaCmd) error {
func (ss *sqlStore) getUserScopeQuota(ctx quota.Context, userID int64) (*quota.Map, error) {
r := quota.Map{}
err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
quotas := make([]*quota.Quota, 0)
if err := sess.Table("quota").Where("user_id=? AND org_id=0", userID).Find(&quotas); err != nil {
return err
@@ -111,7 +111,7 @@ func (ss *sqlStore) getUserScopeQuota(ctx quota.Context, userID int64) (*quota.M
func (ss *sqlStore) getOrgScopeQuota(ctx quota.Context, OrgID int64) (*quota.Map, error) {
r := quota.Map{}
err := ss.db.WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *sqlstore.DBSession) error {
quotas := make([]*quota.Quota, 0)
if err := sess.Table("quota").Where("user_id=0 AND org_id=?", OrgID).Find(&quotas); err != nil {
return err
+1 -1
View File
@@ -20,7 +20,7 @@ func TestIntegrationQuotaDataAccess(t *testing.T) {
t.Skip("skipping integration test")
}
ss := db.InitTestDB(t)
ss := db.InitTestReplDB(t)
quotaStore := sqlStore{
db: ss,
}
+5 -5
View File
@@ -36,17 +36,17 @@ type TestApiKey struct {
ServiceAccountID *int64
}
func SetupUserServiceAccount(t *testing.T, db db.DB, cfg *setting.Cfg, testUser TestUser) *user.User {
func SetupUserServiceAccount(t *testing.T, store db.DB, cfg *setting.Cfg, testUser TestUser) *user.User {
role := string(org.RoleViewer)
if testUser.Role != "" {
role = testUser.Role
}
quotaService := quotaimpl.ProvideService(db, cfg)
orgService, err := orgimpl.ProvideService(db, cfg, quotaService)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), cfg)
orgService, err := orgimpl.ProvideService(store, cfg, quotaService)
require.NoError(t, err)
usrSvc, err := userimpl.ProvideService(
db, orgService, cfg, nil, nil, tracing.InitializeTracerForTest(),
store, orgService, cfg, nil, nil, tracing.InitializeTracerForTest(),
quotaService, supportbundlestest.NewFakeBundleService(),
)
require.NoError(t, err)
@@ -112,7 +112,7 @@ func SetupApiKey(t *testing.T, store db.DB, cfg *setting.Cfg, testKey TestApiKey
func SetupUsersServiceAccounts(t *testing.T, sqlStore db.DB, cfg *setting.Cfg, testUsers []TestUser) (orgID int64) {
role := string(org.RoleNone)
quotaService := quotaimpl.ProvideService(sqlStore, cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(sqlStore), cfg)
orgService, err := orgimpl.ProvideService(sqlStore, cfg, quotaService)
require.NoError(t, err)
usrSvc, err := userimpl.ProvideService(
+11
View File
@@ -264,3 +264,14 @@ func newReplStore(primary *SQLStore, readReplicas ...*SQLStore) *ReplStore {
ret.repls = readReplicas
return ret
}
// FakeReplStoreFromStore returns a ReplStore with the given primary
// SQLStore and no read replicas. This is a bare-minimum wrapper for testing,
// and should be removed when all services are using ReplStore in favor of
// InitTestReplDB.
func FakeReplStoreFromStore(primary *SQLStore) *ReplStore {
return &ReplStore{
SQLStore: primary,
next: 0,
}
}
+4 -3
View File
@@ -17,6 +17,7 @@ import (
"github.com/grafana/grafana/pkg/services/org/orgimpl"
"github.com/grafana/grafana/pkg/services/quota/quotaimpl"
"github.com/grafana/grafana/pkg/services/serviceaccounts"
"github.com/grafana/grafana/pkg/services/sqlstore"
"github.com/grafana/grafana/pkg/services/supportbundles/supportbundlestest"
"github.com/grafana/grafana/pkg/services/team"
"github.com/grafana/grafana/pkg/services/team/sortopts"
@@ -48,7 +49,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) {
},
},
}
quotaService := quotaimpl.ProvideService(sqlStore, cfg)
quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(sqlStore), cfg)
orgSvc, err := orgimpl.ProvideService(sqlStore, cfg, quotaService)
require.NoError(t, err)
userSvc, err := userimpl.ProvideService(
@@ -435,7 +436,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) {
t.Run("Should be able to exclude service accounts from teamembers", func(t *testing.T) {
sqlStore = db.InitTestDB(t)
quotaService := quotaimpl.ProvideService(sqlStore, cfg)
quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(sqlStore), cfg)
orgSvc, err := orgimpl.ProvideService(sqlStore, cfg, quotaService)
require.NoError(t, err)
userSvc, err := userimpl.ProvideService(
@@ -573,7 +574,7 @@ func TestIntegrationSQLStore_GetTeamMembers_ACFilter(t *testing.T) {
require.NoError(t, errCreateTeam)
team2, errCreateTeam := teamSvc.CreateTeam(context.Background(), "group2 name", "test2@example.org", testOrgID)
require.NoError(t, errCreateTeam)
quotaService := quotaimpl.ProvideService(store, cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), cfg)
orgSvc, err := orgimpl.ProvideService(store, cfg, quotaService)
require.NoError(t, err)
userSvc, err := userimpl.ProvideService(
+3 -3
View File
@@ -34,7 +34,7 @@ func TestIntegrationUserDataAccess(t *testing.T) {
}
ss, cfg := db.InitTestDBWithCfg(t)
quotaService := quotaimpl.ProvideService(ss, cfg)
quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(ss), cfg)
orgService, err := orgimpl.ProvideService(ss, cfg, quotaService)
require.NoError(t, err)
userStore := ProvideStore(ss, setting.NewCfg())
@@ -903,7 +903,7 @@ func createFiveTestUsers(t *testing.T, svc user.Service, fn func(i int) *user.Cr
func TestMetricsUsage(t *testing.T) {
ss, cfg := db.InitTestDBWithCfg(t)
userStore := ProvideStore(ss, setting.NewCfg())
quotaService := quotaimpl.ProvideService(ss, cfg)
quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(ss), cfg)
orgService, err := orgimpl.ProvideService(ss, cfg, quotaService)
require.NoError(t, err)
@@ -962,7 +962,7 @@ func assertEqualUser(t *testing.T, expected, got *user.User) {
func createOrgAndUserSvc(t *testing.T, store db.DB, cfg *setting.Cfg) (org.Service, user.Service) {
t.Helper()
quotaService := quotaimpl.ProvideService(store, cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), cfg)
orgService, err := orgimpl.ProvideService(store, cfg, quotaService)
require.NoError(t, err)
usrSvc, err := ProvideService(