IntenseWebs/grafana
3
0
Fork 0
mirror of https://github.com/grafana/grafana.git synced 2024-11-21 16:38:03 -06:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Kristin Laemmert 2024-08-08 13:41:33 -04:00 committed by GitHub
parent 787abccfbc
commit 299c142f6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
20 changed files with 74 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pkg/api/org_users_test.go
View File

@ -41,7 +41,7 @@ func setUpGetOrgUsersDB(t *testing.T, sqlStore db.DB, cfg *setting.Cfg) {
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = int(testOrgID)
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(

2
pkg/cmd/grafana-cli/commands/conflict_user_command.go
View File

@ -79,7 +79,7 @@ func initializeConflictResolver(cmd *utils.ContextCommandLine, f Formatter, ctx
if err != nil {
return nil, fmt.Errorf("%v: %w", "failed to get users with conflicting logins", err)
}
quotaService := quotaimpl.ProvideService(s, cfg)
quotaService := quotaimpl.ProvideService(replstore, cfg)
userService, err := userimpl.ProvideService(s, nil, cfg, nil, nil, tracer, quotaService, supportbundlestest.NewFakeBundleService())
if err != nil {
return nil, fmt.Errorf("%v: %w", "failed to get user service", err)

6
pkg/infra/db/dbrepl.go
View File

@ -11,3 +11,9 @@ type ReplDB interface {
// ReadReplica is the read-only database connection. If no read replica is configured, the implementation must return the primary DB.
ReadReplica() *sqlstore.SQLStore
}
// FakeREplDBFromDBForTests returns a ReplDB that uses the given DB as the primary connection. It's a helper function for tests.
func FakeReplDBFromDB(primary DB) ReplDB {
ss := primary.(*sqlstore.SQLStore)
return sqlstore.FakeReplStoreFromStore(ss)
}

2
pkg/services/org/orgimpl/store_test.go
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(

2
pkg/services/quota/quotaimpl/quota.go
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},

10
pkg/services/quota/quotaimpl/store.go
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

2
pkg/services/quota/quotaimpl/store_test.go
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,
}

10
pkg/services/serviceaccounts/tests/common.go
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
pkg/services/sqlstore/replstore.go
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,
}
}

7
pkg/services/team/teamimpl/store_test.go
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(

6
pkg/services/user/userimpl/store_test.go
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(

8
pkg/tests/api/alerting/api_alertmanager_test.go
View File

@ -2656,17 +2656,17 @@ func rulesNamespaceWithoutVariableValues(t *testing.T, b []byte) (string, map[st
return string(json), m
}
func createUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
func createUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = 1
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)

5
pkg/tests/api/correlations/common_test.go
View File

@ -9,6 +9,7 @@ import (
"github.com/stretchr/testify/require"
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/server"
"github.com/grafana/grafana/pkg/services/correlations"
@ -144,7 +145,7 @@ func (c TestContext) createOrg(name string) int64 {
c.t.Helper()
store := c.env.SQLStore
c.env.Cfg.AutoAssignOrg = false
quotaService := quotaimpl.ProvideService(store, c.env.Cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), c.env.Cfg)
orgService, err := orgimpl.ProvideService(store, c.env.Cfg, quotaService)
require.NoError(c.t, err)
orgId, err := orgService.GetOrCreate(context.Background(), name)
@ -158,7 +159,7 @@ func (c TestContext) createUser(cmd user.CreateUserCommand) User {
c.env.Cfg.AutoAssignOrg = true
c.env.Cfg.AutoAssignOrgId = 1
quotaService := quotaimpl.ProvideService(store, c.env.Cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), c.env.Cfg)
orgService, err := orgimpl.ProvideService(store, c.env.Cfg, quotaService)
require.NoError(c.t, err)
usrSvc, err := userimpl.ProvideService(

8
pkg/tests/api/dashboards/api_dashboards_test.go
View File

@ -115,17 +115,17 @@ func TestIntegrationDashboardQuota(t *testing.T) {
})
}
func createUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
func createUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = 1
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)

13
pkg/tests/api/folders/api_folder_test.go
View File

@ -9,6 +9,9 @@ import (
"github.com/go-openapi/runtime"
"github.com/grafana/grafana-openapi-client-go/client/folders"
"github.com/grafana/grafana-openapi-client-go/models"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/services/featuremgmt"
@ -21,8 +24,6 @@ import (
"github.com/grafana/grafana/pkg/setting"
"github.com/grafana/grafana/pkg/tests"
"github.com/grafana/grafana/pkg/tests/testinfra"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const orgID = 1
@ -207,17 +208,17 @@ func TestIntegrationNestedFoldersOn(t *testing.T) {
})
}
func createUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
func createUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = orgID
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)

8
pkg/tests/api/plugins/api_plugins_test.go
View File

@ -193,17 +193,17 @@ func TestIntegrationPluginAssets(t *testing.T) {
})
}
func createUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) {
func createUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = 1
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)

8
pkg/tests/api/stats/admin_test.go
View File

@ -81,17 +81,17 @@ func grafanaSetup(t *testing.T, opts testinfra.GrafanaOpts) string {
return fmt.Sprintf("http://%s:%s@%s/api/admin/stats", "grafana", "password", grafanaListedAddr)
}
func createUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
func createUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = 1
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)

4
pkg/tests/apis/helper.go
View File

@ -19,13 +19,13 @@ import (
"k8s.io/apimachinery/pkg/runtime/schema"
"k8s.io/apimachinery/pkg/runtime/serializer/yaml"
yamlutil "k8s.io/apimachinery/pkg/util/yaml"
"k8s.io/client-go/discovery"
"k8s.io/client-go/dynamic"
"k8s.io/client-go/rest"
"github.com/grafana/grafana/pkg/apimachinery/identity"
"github.com/grafana/grafana/pkg/apimachinery/utils"
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/infra/localcache"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/server"
@ -399,7 +399,7 @@ func (c *K8sTestHelper) CreateUser(name string, orgName string, basicRole org.Ro
c.env.Cfg.AutoAssignOrgId = 1 // the default
}()
quotaService := quotaimpl.ProvideService(store, c.env.Cfg)
quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), c.env.Cfg)
orgService, err := orgimpl.ProvideService(store, c.env.Cfg, quotaService)
require.NoError(c.t, err)

2
pkg/tests/testinfra/testinfra.go
View File

@ -451,7 +451,7 @@ func CreateUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUser
cfg.AutoAssignOrgId = 1
cmd.OrgID = 1
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 := userimpl.ProvideService(

11
pkg/tests/utils.go
View File

@ -9,6 +9,8 @@ import (
"github.com/go-openapi/strfmt"
goapi "github.com/grafana/grafana-openapi-client-go/client"
"github.com/stretchr/testify/require"
"github.com/grafana/grafana/pkg/infra/db"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/services/accesscontrol/resourcepermissions"
@ -19,20 +21,19 @@ import (
"github.com/grafana/grafana/pkg/services/user"
"github.com/grafana/grafana/pkg/services/user/userimpl"
"github.com/grafana/grafana/pkg/setting"
"github.com/stretchr/testify/require"
)
func CreateUser(t *testing.T, db db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
func CreateUser(t *testing.T, store db.DB, cfg *setting.Cfg, cmd user.CreateUserCommand) int64 {
t.Helper()
cfg.AutoAssignOrg = true
cfg.AutoAssignOrgId = 1
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)