session: add a CheckSessionGroupPredicate method

This method can be used to check that each session in a group passes for
a given predicate.

Author: ellemouton

session/interface.go

@@ -171,5 +171,11 @@ type Store interface {
 	// GetSessionByID fetches the session with the given ID.
 	GetSessionByID(id ID) (*Session, error)
 
+	// CheckSessionGroupPredicate iterates over all the sessions in a group
+	// and checks if each one passes the given predicate function. True is
+	// returned if each session passes.
+	CheckSessionGroupPredicate(groupID ID,
+		fn func(s *Session) bool) (bool, error)
+
 	IDToGroupIndex
 }

session/store.go

@@ -385,39 +385,115 @@ func (db *DB) GetGroupID(sessionID ID) (ID, error) {
 //
 // NOTE: this is part of the IDToGroupIndex interface.
 func (db *DB) GetSessionIDs(groupID ID) ([]ID, error) {
-	var sessionIDs []ID
+	var (
+		sessionIDs []ID
+		err        error
+	)
+	err = db.View(func(tx *bbolt.Tx) error {
+		sessionIDs, err = getSessionIDs(tx, groupID)
+
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return sessionIDs, nil
+}
+
+// CheckSessionGroupPredicate iterates over all the sessions in a group and
+// checks if each one passes the given predicate function. True is returned if
+// each session passes.
+//
+// NOTE: this is part of the Store interface.
+func (db *DB) CheckSessionGroupPredicate(groupID ID,
+	fn func(s *Session) bool) (bool, error) {
+
+	var (
+		pass          bool
+		errFailedPred = errors.New("session failed predicate")
+	)
 	err := db.View(func(tx *bbolt.Tx) error {
 		sessionBkt, err := getBucket(tx, sessionBucketKey)
 		if err != nil {
 			return err
 		}
 
-		groupIndexBkt := sessionBkt.Bucket(groupIDIndexKey)
-		if groupIndexBkt == nil {
-			return ErrDBInitErr
+		sessionIDs, err := getSessionIDs(tx, groupID)
+		if err != nil {
+			return err
 		}
 
-		groupIDBkt := groupIndexBkt.Bucket(groupID[:])
-		if groupIDBkt == nil {
-			return fmt.Errorf("no sessions for group ID %v",
-				groupID)
-		}
+		// Iterate over all the sessions.
+		for _, id := range sessionIDs {
+			key, err := getKeyForID(sessionBkt, id)
+			if err != nil {
+				return err
+			}
+
+			v := sessionBkt.Get(key)
+			if len(v) == 0 {
+				return ErrSessionNotFound
+			}
+
+			session, err := DeserializeSession(bytes.NewReader(v))
+			if err != nil {
+				return err
+			}
 
-		sessionIDsBkt := groupIDBkt.Bucket(sessionIDKey)
-		if sessionIDsBkt == nil {
-			return fmt.Errorf("no sessions for group ID %v",
-				groupID)
+			if !fn(session) {
+				return errFailedPred
+			}
 		}
 
-		return sessionIDsBkt.ForEach(func(_,
-			sessionIDBytes []byte) error {
+		pass = true
 
-			var sessionID ID
-			copy(sessionID[:], sessionIDBytes)
-			sessionIDs = append(sessionIDs, sessionID)
+		return nil
+	})
+	if errors.Is(err, errFailedPred) {
+		return pass, nil
+	}
+	if err != nil {
+		return pass, err
+	}
 
-			return nil
-		})
+	return pass, nil
+}
+
+// getSessionIDs returns all the session IDs associated with the given group ID.
+func getSessionIDs(tx *bbolt.Tx, groupID ID) ([]ID, error) {
+	var sessionIDs []ID
+
+	sessionBkt, err := getBucket(tx, sessionBucketKey)
+	if err != nil {
+		return nil, err
+	}
+
+	groupIndexBkt := sessionBkt.Bucket(groupIDIndexKey)
+	if groupIndexBkt == nil {
+		return nil, ErrDBInitErr
+	}
+
+	groupIDBkt := groupIndexBkt.Bucket(groupID[:])
+	if groupIDBkt == nil {
+		return nil, fmt.Errorf("no sessions for group ID %v",
+			groupID)
+	}
+
+	sessionIDsBkt := groupIDBkt.Bucket(sessionIDKey)
+	if sessionIDsBkt == nil {
+		return nil, fmt.Errorf("no sessions for group ID %v",
+			groupID)
+	}
+
+	err = sessionIDsBkt.ForEach(func(_,
+		sessionIDBytes []byte) error {
+
+		var sessionID ID
+		copy(sessionID[:], sessionIDBytes)
+		sessionIDs = append(sessionIDs, sessionID)
+
+		return nil
 	})
 	if err != nil {
 		return nil, err

session/store_test.go

@@ -1,6 +1,7 @@
 package session
 
 import (
+	"strings"
 	"testing"
 	"time"
 
@@ -172,6 +173,94 @@ func TestLinkedSessions(t *testing.T) {
 	require.EqualValues(t, []ID{s4.ID, s5.ID}, sIDs)
 }
 
+// TestCheckSessionGroupPredicate asserts that the CheckSessionGroupPredicate
+// method correctly checks if each session in a group passes a predicate.
+func TestCheckSessionGroupPredicate(t *testing.T) {
+	// Set up a new DB.
+	db, err := NewDB(t.TempDir(), "test.db")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = db.Close()
+	})
+
+	// We will use the Label of the Session to test that the predicate
+	// function is checked correctly.
+
+	// Add a new session to the DB.
+	s1 := newSession(t, db, "label 1", nil)
+	require.NoError(t, db.CreateSession(s1))
+
+	// Check that the group passes against an appropriate predicate.
+	ok, err := db.CheckSessionGroupPredicate(
+		s1.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label 1")
+		},
+	)
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	// Check that the group fails against an appropriate predicate.
+	ok, err = db.CheckSessionGroupPredicate(
+		s1.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label 2")
+		},
+	)
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	// Add a new session to the same group as the first one.
+	s2 := newSession(t, db, "label 2", &s1.GroupID)
+	require.NoError(t, db.CreateSession(s2))
+
+	// Check that the group passes against an appropriate predicate.
+	ok, err = db.CheckSessionGroupPredicate(
+		s1.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label")
+		},
+	)
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	// Check that the group fails against an appropriate predicate.
+	ok, err = db.CheckSessionGroupPredicate(
+		s1.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label 1")
+		},
+	)
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	// Add a new session that is not linked to the first one.
+	s3 := newSession(t, db, "completely different", nil)
+	require.NoError(t, db.CreateSession(s3))
+
+	// Ensure that the first group is unaffected.
+	ok, err = db.CheckSessionGroupPredicate(
+		s1.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label")
+		},
+	)
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	// And that the new session is evaluated separately.
+	ok, err = db.CheckSessionGroupPredicate(
+		s3.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "label")
+		},
+	)
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	ok, err = db.CheckSessionGroupPredicate(
+		s3.GroupID, func(s *Session) bool {
+			return strings.Contains(s.Label, "different")
+		},
+	)
+	require.NoError(t, err)
+	require.True(t, ok)
+}
+
 func newSession(t *testing.T, db Store, label string,
 	linkedGroupID *ID) *Session {