AuthZ: Improve team ID fetching for signedInUser (#78378)

* improve team ID fetching for signedInUser

* remove inner join

* rename func

* nit: remove extra params

* nit: spacing and wrapping

pkg/services/team/model.go

@@ -79,6 +79,11 @@ type GetTeamByIDQuery struct {
 // FilterIgnoreUser is used in a get / search teams query when the caller does not want to filter teams by user ID / membership
 const FilterIgnoreUser int64 = 0
 
+type GetTeamIDsByUserQuery struct {
+	OrgID  int64
+	UserID int64 `json:"userId"`
+}
+
 type GetTeamsByUserQuery struct {
 	OrgID        int64
 	UserID       int64 `json:"userId"`

pkg/services/team/team.go

@@ -13,6 +13,7 @@ type Service interface {
 	SearchTeams(ctx context.Context, query *SearchTeamsQuery) (SearchTeamQueryResult, error)
 	GetTeamByID(ctx context.Context, query *GetTeamByIDQuery) (*TeamDTO, error)
 	GetTeamsByUser(ctx context.Context, query *GetTeamsByUserQuery) ([]*TeamDTO, error)
+	GetTeamIDsByUser(ctx context.Context, query *GetTeamIDsByUserQuery) ([]int64, error)
 	AddTeamMember(userID, orgID, teamID int64, isExternal bool, permission dashboards.PermissionType) error
 	UpdateTeamMember(ctx context.Context, cmd *UpdateTeamMemberCommand) error
 	IsTeamMember(orgId int64, teamId int64, userId int64) (bool, error)

pkg/services/team/teamimpl/store.go

@@ -23,6 +23,7 @@ type store interface {
 	Search(ctx context.Context, query *team.SearchTeamsQuery) (team.SearchTeamQueryResult, error)
 	GetByID(ctx context.Context, query *team.GetTeamByIDQuery) (*team.TeamDTO, error)
 	GetByUser(ctx context.Context, query *team.GetTeamsByUserQuery) ([]*team.TeamDTO, error)
+	GetIDsByUser(ctx context.Context, query *team.GetTeamIDsByUserQuery) ([]int64, error)
 	RemoveUsersMemberships(ctx context.Context, userID int64) error
 	AddMember(userID, orgID, teamID int64, isExternal bool, permission dashboards.PermissionType) error
 	UpdateMember(ctx context.Context, cmd *team.UpdateTeamMemberCommand) error
@@ -336,6 +337,22 @@ func (ss *xormStore) GetByUser(ctx context.Context, query *team.GetTeamsByUserQu
 	return queryResult, nil
 }
 
+// GetIDsByUser returns a list of team IDs for the given user
+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 {
+		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)
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get team IDs by user: %w", err)
+	}
+
+	return queryResult, nil
+}
+
 // AddTeamMember adds a user to a team
 func (ss *xormStore) AddMember(userID, orgID, teamID int64, isExternal bool, permission dashboards.PermissionType) error {
 	return ss.db.WithTransactionalDbSession(context.Background(), func(sess *db.Session) error {

pkg/services/team/teamimpl/store_test.go

@@ -125,6 +125,13 @@ func TestIntegrationTeamCommandsAndQueries(t *testing.T) {
 				require.Equal(t, team1.Email, "test1@test.com")
 				require.Equal(t, team1.OrgID, testOrgID)
 				require.EqualValues(t, team1.MemberCount, 2)
+
+				getIDsQuery := &team.GetTeamIDsByUserQuery{OrgID: testOrgID, UserID: userIds[0]}
+				getIDResult, err := teamSvc.GetTeamIDsByUser(context.Background(), getIDsQuery)
+				require.NoError(t, err)
+
+				require.Equal(t, len(getIDResult), 1)
+				require.Equal(t, getIDResult[0], team1.ID)
 			})
 
 			t.Run("Should return latest auth module for users when getting team members", func(t *testing.T) {

pkg/services/team/teamimpl/team.go

@@ -41,6 +41,10 @@ func (s *Service) GetTeamsByUser(ctx context.Context, query *team.GetTeamsByUser
 	return s.store.GetByUser(ctx, query)
 }
 
+func (s *Service) GetTeamIDsByUser(ctx context.Context, query *team.GetTeamIDsByUserQuery) ([]int64, error) {
+	return s.store.GetIDsByUser(ctx, query)
+}
+
 func (s *Service) AddTeamMember(userID, orgID, teamID int64, isExternal bool, permission dashboards.PermissionType) error {
 	return s.store.AddMember(userID, orgID, teamID, isExternal, permission)
 }

pkg/services/team/teamtest/team.go

@@ -75,3 +75,12 @@ func (s *FakeService) GetTeamMembers(ctx context.Context, query *team.GetTeamMem
 
 func (s *FakeService) RegisterDelete(query string) {
 }
+
+func (s *FakeService) GetTeamIDsByUser(ctx context.Context, query *team.GetTeamIDsByUserQuery) ([]int64, error) {
+	result := make([]int64, 0)
+	for _, team := range s.ExpectedTeamsByUser {
+		result = append(result, team.ID)
+	}
+
+	return result, s.ExpectedError
+}

pkg/services/user/userimpl/user.go

@@ -312,29 +312,15 @@ func (s *Service) GetSignedInUser(ctx context.Context, query *user.GetSignedInUs
 		return nil, err
 	}
 
-	// tempUser is used to retrieve the teams for the signed in user for internal use.
+	getTeamsByUserQuery := &team.GetTeamIDsByUserQuery{
-	tempUser := &user.SignedInUser{
+		OrgID:  signedInUser.OrgID,
-		OrgID: signedInUser.OrgID,
+		UserID: signedInUser.UserID,
-		Permissions: map[int64]map[string][]string{
-			signedInUser.OrgID: {
-				ac.ActionTeamsRead: {ac.ScopeTeamsAll},
-			},
-		},
 	}
-	getTeamsByUserQuery := &team.GetTeamsByUserQuery{
+	signedInUser.Teams, err = s.teamService.GetTeamIDsByUser(ctx, getTeamsByUserQuery)
-		OrgID:        signedInUser.OrgID,
-		UserID:       signedInUser.UserID,
-		SignedInUser: tempUser,
-	}
-	getTeamsByUserQueryResult, err := s.teamService.GetTeamsByUser(ctx, getTeamsByUserQuery)
 	if err != nil {
 		return nil, err
 	}
 
-	signedInUser.Teams = make([]int64, len(getTeamsByUserQueryResult))
-	for i, t := range getTeamsByUserQueryResult {
-		signedInUser.Teams[i] = t.ID
-	}
 	return signedInUser, err
 }