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