core/grpc: add IterateAll method

This commit is contained in:
Caleb Doxsey 2024-08-19 16:51:23 -06:00
parent 2925447d35
commit ae4d266fa4
3 changed files with 198 additions and 18 deletions

View file

@ -8,43 +8,32 @@ import (
"github.com/pomerium/pomerium/internal/sets"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/grpc/session"
"github.com/pomerium/pomerium/pkg/protoutil"
)
var sessionTypeURL = protoutil.GetTypeURL(new(session.Session))
// CurrentUsers returns a list of users active within the current UTC day
func CurrentUsers(
ctx context.Context,
client databroker.DataBrokerServiceClient,
) ([]string, error) {
records, _, _, err := databroker.InitialSync(ctx, client, &databroker.SyncLatestRequest{
Type: sessionTypeURL,
})
if err != nil {
return nil, fmt.Errorf("fetching sessions: %w", err)
}
users := sets.NewHash[string]()
utcNow := time.Now().UTC()
threshold := time.Date(utcNow.Year(), utcNow.Month(), utcNow.Day(), 0, 0, 0, 0, time.UTC)
for _, record := range records {
var s session.Session
err := record.GetData().UnmarshalTo(&s)
for s, err := range databroker.IterateAll[session.Session](ctx, client) {
if err != nil {
return nil, fmt.Errorf("unmarshaling session: %w", err)
return nil, fmt.Errorf("error fetching sessions: %w", err)
}
if s.UserId == "" { // session creation is in progress
if s.Object.GetUserId() == "" { // session creation is in progress
continue
}
if s.AccessedAt == nil {
if s.Object.GetAccessedAt() == nil {
continue
}
if s.AccessedAt.AsTime().Before(threshold) {
if s.Object.GetAccessedAt().AsTime().Before(threshold) {
continue
}
users.Add(s.UserId)
users.Add(s.Object.GetUserId())
}
return users.Items(), nil