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
pull/91706/head
Kristin Laemmert 1 year ago committed by GitHub
parent 787abccfbc
commit 299c142f6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 2
      pkg/api/org_users_test.go
  2. 2
      pkg/cmd/grafana-cli/commands/conflict_user_command.go
  3. 6
      pkg/infra/db/dbrepl.go
  4. 2
      pkg/services/org/orgimpl/store_test.go
  5. 2
      pkg/services/quota/quotaimpl/quota.go
  6. 10
      pkg/services/quota/quotaimpl/store.go
  7. 2
      pkg/services/quota/quotaimpl/store_test.go
  8. 10
      pkg/services/serviceaccounts/tests/common.go
  9. 11
      pkg/services/sqlstore/replstore.go
  10. 7
      pkg/services/team/teamimpl/store_test.go
  11. 6
      pkg/services/user/userimpl/store_test.go
  12. 8
      pkg/tests/api/alerting/api_alertmanager_test.go
  13. 5
      pkg/tests/api/correlations/common_test.go
  14. 8
      pkg/tests/api/dashboards/api_dashboards_test.go
  15. 13
      pkg/tests/api/folders/api_folder_test.go
  16. 8
      pkg/tests/api/plugins/api_plugins_test.go
  17. 8
      pkg/tests/api/stats/admin_test.go
  18. 4
      pkg/tests/apis/helper.go
  19. 2
      pkg/tests/testinfra/testinfra.go
  20. 11
      pkg/tests/utils.go

@ -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(

@ -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)

@ -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)
}

@ -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(

@ -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},

@ -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

@ -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,
}

@ -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(

@ -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,
}
}

@ -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(

@ -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(

@ -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)

@ -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(

@ -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)

@ -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)

@ -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)

@ -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)

@ -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)

@ -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(

@ -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)

Loading…
Cancel
Save