From c6ae2d7999aa6fc797db39e9d66c6fea70278f83 Mon Sep 17 00:00:00 2001 From: Kristin Laemmert Date: Thu, 15 Aug 2024 12:19:38 -0400 Subject: [PATCH] chore: add replDB to team service (#91799) --- .../commands/conflict_user_command_test.go | 6 +-- .../resourcepermissions/service_test.go | 2 +- .../resourcepermissions/store_bench_test.go | 2 +- pkg/services/team/teamimpl/store.go | 47 ++++++++++--------- pkg/services/team/teamimpl/store_test.go | 32 ++++++------- pkg/services/team/teamimpl/team.go | 2 +- pkg/tests/apis/helper.go | 6 +-- 7 files changed, 49 insertions(+), 48 deletions(-) diff --git a/pkg/cmd/grafana-cli/commands/conflict_user_command_test.go b/pkg/cmd/grafana-cli/commands/conflict_user_command_test.go index c0aafacd3c0..a43a07e1e53 100644 --- a/pkg/cmd/grafana-cli/commands/conflict_user_command_test.go +++ b/pkg/cmd/grafana-cli/commands/conflict_user_command_test.go @@ -636,7 +636,7 @@ func TestIntegrationMergeUser(t *testing.T) { } t.Run("should be able to merge user", func(t *testing.T) { // Restore after destructive operation - sqlStore := db.InitTestDB(t) + sqlStore := db.InitTestReplDB(t) teamSvc, err := teamimpl.ProvideService(sqlStore, setting.NewCfg(), tracing.InitializeTracerForTest()) require.NoError(t, err) team1, err := teamSvc.CreateTeam(context.Background(), "team1 name", "", 1) @@ -714,10 +714,10 @@ func TestIntegrationMergeUser(t *testing.T) { } // get users - conflictUsers, err := GetUsersWithConflictingEmailsOrLogins(&cli.Context{Context: context.Background()}, sqlStore) + conflictUsers, err := GetUsersWithConflictingEmailsOrLogins(&cli.Context{Context: context.Background()}, sqlStore.SQLStore) require.NoError(t, err) r := ConflictResolver{ - Store: sqlStore, + Store: sqlStore.SQLStore, userService: usertest.NewUserServiceFake(), ac: actest.FakeService{}, } diff --git a/pkg/services/accesscontrol/resourcepermissions/service_test.go b/pkg/services/accesscontrol/resourcepermissions/service_test.go index 94ead810dd7..c8c9dbb406c 100644 --- a/pkg/services/accesscontrol/resourcepermissions/service_test.go +++ b/pkg/services/accesscontrol/resourcepermissions/service_test.go @@ -493,7 +493,7 @@ func setupTestEnvironment(t *testing.T, ops Options) (*Service, user.Service, te cfg := setting.NewCfg() tracer := tracing.InitializeTracerForTest() - teamSvc, err := teamimpl.ProvideService(sql, cfg, tracer) + teamSvc, err := teamimpl.ProvideService(db.FakeReplDBFromDB(sql), cfg, tracer) require.NoError(t, err) orgSvc, err := orgimpl.ProvideService(sql, cfg, quotatest.New(false, nil)) diff --git a/pkg/services/accesscontrol/resourcepermissions/store_bench_test.go b/pkg/services/accesscontrol/resourcepermissions/store_bench_test.go index 51a63a5deba..60addeeee24 100644 --- a/pkg/services/accesscontrol/resourcepermissions/store_bench_test.go +++ b/pkg/services/accesscontrol/resourcepermissions/store_bench_test.go @@ -140,7 +140,7 @@ func GenerateDatasourcePermissions(b *testing.B, db db.DB, cfg *setting.Cfg, ac } func generateTeamsAndUsers(b *testing.B, store db.DB, cfg *setting.Cfg, users int) ([]int64, []int64) { - teamSvc, err := teamimpl.ProvideService(store, cfg, tracing.InitializeTracerForTest()) + teamSvc, err := teamimpl.ProvideService(db.FakeReplDBFromDB(store), cfg, tracing.InitializeTracerForTest()) require.NoError(b, err) numberOfTeams := int(math.Ceil(float64(users) / UsersPerTeam)) globalUserId := 0 diff --git a/pkg/services/team/teamimpl/store.go b/pkg/services/team/teamimpl/store.go index 8f3db689fd3..d1bc2afa06b 100644 --- a/pkg/services/team/teamimpl/store.go +++ b/pkg/services/team/teamimpl/store.go @@ -32,7 +32,7 @@ type store interface { } type xormStore struct { - db db.DB + db db.ReplDB cfg *setting.Cfg deletes []string } @@ -85,7 +85,7 @@ func (ss *xormStore) Create(name, email string, orgID int64) (team.Team, error) Created: time.Now(), Updated: time.Now(), } - err := ss.db.WithTransactionalDbSession(context.Background(), func(sess *db.Session) error { + err := ss.db.DB().WithTransactionalDbSession(context.Background(), func(sess *db.Session) error { if isNameTaken, err := isTeamNameTaken(orgID, name, 0, sess); err != nil { return err } else if isNameTaken { @@ -99,7 +99,7 @@ func (ss *xormStore) Create(name, email string, orgID int64) (team.Team, error) } func (ss *xormStore) Update(ctx context.Context, cmd *team.UpdateTeamCommand) error { - return ss.db.WithTransactionalDbSession(ctx, func(sess *db.Session) error { + return ss.db.DB().WithTransactionalDbSession(ctx, func(sess *db.Session) error { if isNameTaken, err := isTeamNameTaken(cmd.OrgID, cmd.Name, cmd.ID, sess); err != nil { return err } else if isNameTaken { @@ -130,7 +130,7 @@ func (ss *xormStore) Update(ctx context.Context, cmd *team.UpdateTeamCommand) er // DeleteTeam will delete a team, its member and any permissions connected to the team func (ss *xormStore) Delete(ctx context.Context, cmd *team.DeleteTeamCommand) error { - return ss.db.WithTransactionalDbSession(ctx, func(sess *db.Session) error { + return ss.db.DB().WithTransactionalDbSession(ctx, func(sess *db.Session) error { if _, err := teamExists(cmd.OrgID, cmd.ID, sess); err != nil { return err } @@ -181,7 +181,7 @@ func (ss *xormStore) Search(ctx context.Context, query *team.SearchTeamsQuery) ( queryResult := team.SearchTeamQueryResult{ Teams: make([]*team.TeamDTO, 0), } - err := ss.db.WithDbSession(ctx, func(sess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *db.Session) error { queryWithWildcards := "%" + query.Query + "%" var sql bytes.Buffer @@ -192,12 +192,12 @@ func (ss *xormStore) Search(ctx context.Context, query *team.SearchTeamsQuery) ( params = append(params, user) } - sql.WriteString(getTeamSelectSQLBase(ss.db, filteredUsers)) + sql.WriteString(getTeamSelectSQLBase(ss.db.ReadReplica(), filteredUsers)) sql.WriteString(` WHERE team.org_id = ?`) params = append(params, query.OrgID) if query.Query != "" { - sql.WriteString(` and team.name ` + ss.db.GetDialect().LikeStr() + ` ?`) + sql.WriteString(` and team.name ` + ss.db.ReadReplica().GetDialect().LikeStr() + ` ?`) params = append(params, queryWithWildcards) } @@ -234,7 +234,7 @@ func (ss *xormStore) Search(ctx context.Context, query *team.SearchTeamsQuery) ( if query.Limit != 0 { offset := query.Limit * (query.Page - 1) - sql.WriteString(ss.db.GetDialect().LimitOffset(int64(query.Limit), int64(offset))) + sql.WriteString(ss.db.ReadReplica().GetDialect().LimitOffset(int64(query.Limit), int64(offset))) } if err := sess.SQL(sql.String(), params...).Find(&queryResult.Teams); err != nil { @@ -246,7 +246,7 @@ func (ss *xormStore) Search(ctx context.Context, query *team.SearchTeamsQuery) ( countSess.Where("team.org_id=?", query.OrgID) if query.Query != "" { - countSess.Where(`name `+ss.db.GetDialect().LikeStr()+` ?`, queryWithWildcards) + countSess.Where(`name `+ss.db.ReadReplica().GetDialect().LikeStr()+` ?`, queryWithWildcards) } if query.Name != "" { @@ -269,12 +269,12 @@ func (ss *xormStore) Search(ctx context.Context, query *team.SearchTeamsQuery) ( func (ss *xormStore) GetByID(ctx context.Context, query *team.GetTeamByIDQuery) (*team.TeamDTO, error) { var queryResult *team.TeamDTO - err := ss.db.WithDbSession(ctx, func(sess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *db.Session) error { var sql bytes.Buffer params := make([]any, 0) filteredUsers := getFilteredUsers(query.SignedInUser, query.HiddenUsers) - sql.WriteString(getTeamSelectSQLBase(ss.db, filteredUsers)) + sql.WriteString(getTeamSelectSQLBase(ss.db.ReadReplica(), filteredUsers)) for _, user := range filteredUsers { params = append(params, user) } @@ -305,12 +305,12 @@ func (ss *xormStore) GetByID(ctx context.Context, query *team.GetTeamByIDQuery) // GetTeamsByUser is used by the Guardian when checking a users' permissions func (ss *xormStore) GetByUser(ctx context.Context, query *team.GetTeamsByUserQuery) ([]*team.TeamDTO, error) { queryResult := make([]*team.TeamDTO, 0) - err := ss.db.WithDbSession(ctx, func(sess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *db.Session) error { var sql bytes.Buffer var params []any params = append(params, query.OrgID, query.UserID) - sql.WriteString(getTeamSelectSQLBase(ss.db, []string{})) + sql.WriteString(getTeamSelectSQLBase(ss.db.ReadReplica(), []string{})) sql.WriteString(` INNER JOIN team_member on team.id = team_member.team_id`) sql.WriteString(` WHERE team.org_id = ? and team_member.user_id = ?`) @@ -334,7 +334,7 @@ func (ss *xormStore) GetByUser(ctx context.Context, query *team.GetTeamsByUserQu func (ss *xormStore) GetIDsByUser(ctx context.Context, query *team.GetTeamIDsByUserQuery) ([]int64, error) { queryResult := make([]int64, 0) - err := ss.db.WithDbSession(ctx, func(sess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(ctx, func(sess *db.Session) error { return sess.SQL(`SELECT tm.team_id FROM team_member as tm WHERE tm.user_id=? AND tm.org_id=?;`, query.UserID, query.OrgID).Find(&queryResult) @@ -364,7 +364,7 @@ func getTeamMember(sess *db.Session, orgId int64, teamId int64, userId int64) (t func (ss *xormStore) IsMember(orgId int64, teamId int64, userId int64) (bool, error) { var isMember bool - err := ss.db.WithDbSession(context.Background(), func(sess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(context.Background(), func(sess *db.Session) error { var err error isMember, err = isTeamMember(sess, orgId, teamId, userId) return err @@ -461,7 +461,7 @@ func removeTeamMember(sess *db.Session, cmd *team.RemoveTeamMemberCommand) error // RemoveUsersMemberships removes all the team membership entries for the given user. // Only used when removing a user from a Grafana instance. func (ss *xormStore) RemoveUsersMemberships(ctx context.Context, userID int64) error { - return ss.db.WithTransactionalDbSession(ctx, func(sess *db.Session) error { + return ss.db.DB().WithTransactionalDbSession(ctx, func(sess *db.Session) error { var rawSQL = "DELETE FROM team_member WHERE user_id = ?" _, err := sess.Exec(rawSQL, userID) return err @@ -489,7 +489,7 @@ func (ss *xormStore) GetMembers(ctx context.Context, query *team.GetTeamMembersQ // With accesscontrol we filter out users based on the SignedInUser's permissions // Note we assume that checking SignedInUser is allowed to see team members for this team has already been performed // If the signed in user is not set no member will be returned - sqlID := fmt.Sprintf("%s.%s", ss.db.GetDialect().Quote("user"), ss.db.GetDialect().Quote("id")) + sqlID := fmt.Sprintf("%s.%s", ss.db.ReadReplica().GetDialect().Quote("user"), ss.db.ReadReplica().GetDialect().Quote("id")) *acFilter, err = ac.Filter(query.SignedInUser, sqlID, "users:id:", ac.ActionOrgUsersRead) if err != nil { return nil, err @@ -501,15 +501,16 @@ func (ss *xormStore) GetMembers(ctx context.Context, query *team.GetTeamMembersQ // getTeamMembers return a list of members for the specified team func (ss *xormStore) getTeamMembers(ctx context.Context, query *team.GetTeamMembersQuery, acUserFilter *ac.SQLFilter) ([]*team.TeamMemberDTO, error) { queryResult := make([]*team.TeamMemberDTO, 0) - err := ss.db.WithDbSession(ctx, func(dbSess *db.Session) error { + err := ss.db.ReadReplica().WithDbSession(ctx, func(dbSess *db.Session) error { + dialect := ss.db.ReadReplica().GetDialect() sess := dbSess.Table("team_member") - sess.Join("INNER", ss.db.GetDialect().Quote("user"), - fmt.Sprintf("team_member.user_id=%s.%s", ss.db.GetDialect().Quote("user"), ss.db.GetDialect().Quote("id")), + sess.Join("INNER", dialect.Quote("user"), + fmt.Sprintf("team_member.user_id=%s.%s", dialect.Quote("user"), dialect.Quote("id")), ) sess.Join("INNER", "team", "team.id=team_member.team_id") // explicitly check for serviceaccounts - sess.Where(fmt.Sprintf("%s.is_service_account=?", ss.db.GetDialect().Quote("user")), ss.db.GetDialect().BooleanStr(false)) + sess.Where(fmt.Sprintf("%s.is_service_account=?", dialect.Quote("user")), dialect.BooleanStr(false)) if acUserFilter != nil { sess.Where(acUserFilter.Where, acUserFilter.Args...) @@ -521,7 +522,7 @@ func (ss *xormStore) getTeamMembers(ctx context.Context, query *team.GetTeamMemb FROM user_auth WHERE user_auth.user_id = team_member.user_id ORDER BY user_auth.created DESC ` + - ss.db.GetDialect().Limit(1) + ")" + dialect.Limit(1) + ")" sess.Join("LEFT", "user_auth", authJoinCondition) if query.OrgID != 0 { @@ -537,7 +538,7 @@ func (ss *xormStore) getTeamMembers(ctx context.Context, query *team.GetTeamMemb sess.Where("team_member.user_id=?", query.UserID) } if query.External { - sess.Where("team_member.external=?", ss.db.GetDialect().BooleanStr(true)) + sess.Where("team_member.external=?", dialect.BooleanStr(true)) } sess.Cols( "team_member.org_id", diff --git a/pkg/services/team/teamimpl/store_test.go b/pkg/services/team/teamimpl/store_test.go index d198073966e..d3c4bf5cea2 100644 --- a/pkg/services/team/teamimpl/store_test.go +++ b/pkg/services/team/teamimpl/store_test.go @@ -36,7 +36,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { t.Skip("skipping integration test") } t.Run("Testing Team commands and queries", func(t *testing.T) { - sqlStore, cfg := db.InitTestDBWithCfg(t) + sqlStore, cfg := db.InitTestReplDBWithCfg(t) teamSvc, err := ProvideService(sqlStore, cfg, tracing.InitializeTracerForTest()) require.NoError(t, err) testUser := &user.SignedInUser{ @@ -49,7 +49,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }, }, } - quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(sqlStore), cfg) + quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(sqlStore.SQLStore), cfg) orgSvc, err := orgimpl.ProvideService(sqlStore, cfg, quotaService) require.NoError(t, err) userSvc, err := userimpl.ProvideService( @@ -149,7 +149,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }) t.Run("Should return latest auth module for users when getting team members", func(t *testing.T) { - sqlStore = db.InitTestDB(t) + sqlStore = db.InitTestReplDB(t) setup() userId := userIds[1] @@ -200,7 +200,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }) t.Run("Should default to member permission level when updating a user with invalid permission level", func(t *testing.T) { - sqlStore = db.InitTestDB(t) + sqlStore = db.InitTestReplDB(t) setup() userID := userIds[0] @@ -314,7 +314,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }) t.Run("Should be able to return all teams a user is member of", func(t *testing.T) { - sqlStore = db.InitTestDB(t) + sqlStore = db.InitTestReplDB(t) setup() groupId := team2.ID err = sqlStore.WithDbSession(context.Background(), func(sess *db.Session) error { @@ -375,7 +375,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }) t.Run("A user should be able to remove the admin permission if there are other admins", func(t *testing.T) { - sqlStore = db.InitTestDB(t) + sqlStore = db.InitTestReplDB(t) setup() err = sqlStore.WithDbSession(context.Background(), func(sess *db.Session) error { @@ -394,7 +394,7 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) { }) t.Run("Should not return hidden users in team member count", func(t *testing.T) { - sqlStore = db.InitTestDB(t) + sqlStore = db.InitTestReplDB(t) setup() signedInUser := &user.SignedInUser{ Login: "loginuser0", @@ -435,8 +435,8 @@ 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.FakeReplStoreFromStore(sqlStore), cfg) + sqlStore = db.InitTestReplDB(t) + quotaService := quotaimpl.ProvideService(sqlstore.FakeReplStoreFromStore(sqlStore.SQLStore), cfg) orgSvc, err := orgimpl.ProvideService(sqlStore, cfg, quotaService) require.NoError(t, err) userSvc, err := userimpl.ProvideService( @@ -530,7 +530,7 @@ func TestIntegrationSQLStore_SearchTeams(t *testing.T) { }, } - store, cfg := db.InitTestDBWithCfg(t, db.InitTestDBOpt{}) + store, cfg := db.InitTestReplDBWithCfg(t, db.InitTestDBOpt{}) teamSvc, err := ProvideService(store, cfg, tracing.InitializeTracerForTest()) require.NoError(t, err) @@ -567,18 +567,18 @@ func TestIntegrationSQLStore_GetTeamMembers_ACFilter(t *testing.T) { userIds := make([]int64, 4) // Seed 2 teams with 2 members - setup := func(store db.DB, cfg *setting.Cfg) { + setup := func(store db.ReplDB, cfg *setting.Cfg) { teamSvc, err := ProvideService(store, cfg, tracing.InitializeTracerForTest()) require.NoError(t, err) team1, errCreateTeam := teamSvc.CreateTeam(context.Background(), "group1 name", "test1@example.org", testOrgID) require.NoError(t, errCreateTeam) team2, errCreateTeam := teamSvc.CreateTeam(context.Background(), "group2 name", "test2@example.org", testOrgID) require.NoError(t, errCreateTeam) - quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), cfg) - orgSvc, err := orgimpl.ProvideService(store, cfg, quotaService) + quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store.DB()), cfg) + orgSvc, err := orgimpl.ProvideService(store.DB(), cfg, quotaService) require.NoError(t, err) userSvc, err := userimpl.ProvideService( - store, orgSvc, cfg, teamSvc, nil, tracing.InitializeTracerForTest(), + store.DB(), orgSvc, cfg, teamSvc, nil, tracing.InitializeTracerForTest(), quotaService, supportbundlestest.NewFakeBundleService(), ) require.NoError(t, err) @@ -594,7 +594,7 @@ func TestIntegrationSQLStore_GetTeamMembers_ACFilter(t *testing.T) { userIds[i] = user.ID } - errAddMembers := store.WithDbSession(context.Background(), func(sess *db.Session) error { + errAddMembers := store.DB().WithDbSession(context.Background(), func(sess *db.Session) error { err := AddOrUpdateTeamMemberHook(sess, userIds[0], testOrgID, team1.ID, false, 0) if err != nil { return err @@ -612,7 +612,7 @@ func TestIntegrationSQLStore_GetTeamMembers_ACFilter(t *testing.T) { require.NoError(t, errAddMembers) } - store, cfg := db.InitTestDBWithCfg(t, db.InitTestDBOpt{}) + store, cfg := db.InitTestReplDBWithCfg(t, db.InitTestDBOpt{}) setup(store, cfg) teamSvc, err := ProvideService(store, cfg, tracing.InitializeTracerForTest()) require.NoError(t, err) diff --git a/pkg/services/team/teamimpl/team.go b/pkg/services/team/teamimpl/team.go index e7c9537b751..14212f22b7f 100644 --- a/pkg/services/team/teamimpl/team.go +++ b/pkg/services/team/teamimpl/team.go @@ -17,7 +17,7 @@ type Service struct { tracer tracing.Tracer } -func ProvideService(db db.DB, cfg *setting.Cfg, tracer tracing.Tracer) (team.Service, error) { +func ProvideService(db db.ReplDB, cfg *setting.Cfg, tracer tracing.Tracer) (team.Service, error) { return &Service{ store: &xormStore{db: db, cfg: cfg, deletes: []string{}}, tracer: tracer, diff --git a/pkg/tests/apis/helper.go b/pkg/tests/apis/helper.go index 0e96f455238..fe9efe26db0 100644 --- a/pkg/tests/apis/helper.go +++ b/pkg/tests/apis/helper.go @@ -25,7 +25,6 @@ import ( "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" @@ -436,12 +435,13 @@ func (c *K8sTestHelper) CreateUser(name string, orgName string, basicRole org.Ro c.t.Helper() store := c.env.SQLStore + replStore := c.env.ReadReplStore defer func() { c.env.Cfg.AutoAssignOrg = false c.env.Cfg.AutoAssignOrgId = 1 // the default }() - quotaService := quotaimpl.ProvideService(db.FakeReplDBFromDB(store), c.env.Cfg) + quotaService := quotaimpl.ProvideService(replStore, c.env.Cfg) orgService, err := orgimpl.ProvideService(store, c.env.Cfg, quotaService) require.NoError(c.t, err) @@ -462,7 +462,7 @@ func (c *K8sTestHelper) CreateUser(name string, orgName string, basicRole org.Ro c.env.Cfg.AutoAssignOrg = true c.env.Cfg.AutoAssignOrgId = int(orgId) - teamSvc, err := teamimpl.ProvideService(store, c.env.Cfg, tracing.InitializeTracerForTest()) + teamSvc, err := teamimpl.ProvideService(replStore, c.env.Cfg, tracing.InitializeTracerForTest()) require.NoError(c.t, err) cache := localcache.ProvideService()