github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datasets/subjectsetbytype.go

package datasets

import (
	core "github.com/authzed/spicedb/pkg/proto/core/v1"
	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
	"github.com/authzed/spicedb/pkg/tuple"
)

// SubjectByTypeSet is a set of SubjectSet's, grouped by their subject types.
type SubjectByTypeSet struct {
	byType map[string]SubjectSet
}

// NewSubjectByTypeSet creates and returns a new SubjectByTypeSet.
func NewSubjectByTypeSet() *SubjectByTypeSet {
	return &SubjectByTypeSet{
		byType: map[string]SubjectSet{},
	}
}

// AddSubjectOf adds the subject found in the given relationship, along with its caveat.
func (s *SubjectByTypeSet) AddSubjectOf(relationship *core.RelationTuple) error {
	return s.AddSubject(relationship.Subject, relationship.Caveat)
}

// AddConcreteSubject adds a non-caveated subject to the set.
func (s *SubjectByTypeSet) AddConcreteSubject(subject *core.ObjectAndRelation) error {
	return s.AddSubject(subject, nil)
}

// AddSubject adds the specified subject to the set.
func (s *SubjectByTypeSet) AddSubject(subject *core.ObjectAndRelation, caveat *core.ContextualizedCaveat) error {
	key := tuple.JoinRelRef(subject.Namespace, subject.Relation)
	if _, ok := s.byType[key]; !ok {
		s.byType[key] = NewSubjectSet()
	}

	return s.byType[key].Add(&v1.FoundSubject{
		SubjectId:        subject.ObjectId,
		CaveatExpression: wrapCaveat(caveat),
	})
}

// ForEachType invokes the handler for each type of ObjectAndRelation found in the set, along
// with all IDs of objects of that type.
func (s *SubjectByTypeSet) ForEachType(handler func(rr *core.RelationReference, subjects SubjectSet)) {
	for key, subjects := range s.byType {
		ns, rel := tuple.MustSplitRelRef(key)
		handler(&core.RelationReference{
			Namespace: ns,
			Relation:  rel,
		}, subjects)
	}
}

// ForEachTypeUntil invokes the handler for each type of ObjectAndRelation found in the set, along
// with all IDs of objects of that type, until the handler returns an error or false.
func (s *SubjectByTypeSet) ForEachTypeUntil(handler func(rr *core.RelationReference, subjects SubjectSet) (bool, error)) error {
	for key, subjects := range s.byType {
		ns, rel := tuple.MustSplitRelRef(key)
		ok, err := handler(&core.RelationReference{
			Namespace: ns,
			Relation:  rel,
		}, subjects)
		if err != nil {
			return err
		}
		if !ok {
			return nil
		}
	}
	return nil
}

// Map runs the mapper function over each type of object in the set, returning a new ONRByTypeSet with
// the object type replaced by that returned by the mapper function.
func (s *SubjectByTypeSet) Map(mapper func(rr *core.RelationReference) (*core.RelationReference, error)) (*SubjectByTypeSet, error) {
	mapped := NewSubjectByTypeSet()
	for key, subjectset := range s.byType {
		ns, rel := tuple.MustSplitRelRef(key)
		updatedType, err := mapper(&core.RelationReference{
			Namespace: ns,
			Relation:  rel,
		})
		if err != nil {
			return nil, err
		}
		if updatedType == nil {
			continue
		}

		key := tuple.JoinRelRef(updatedType.Namespace, updatedType.Relation)
		if existing, ok := mapped.byType[key]; ok {
			cloned := subjectset.Clone()
			if err := cloned.UnionWithSet(existing); err != nil {
				return nil, err
			}
			mapped.byType[key] = cloned
		} else {
			mapped.byType[key] = subjectset
		}
	}
	return mapped, nil
}

// IsEmpty returns true if the set is empty.
func (s *SubjectByTypeSet) IsEmpty() bool {
	return len(s.byType) == 0
}

// Len returns the number of keys in the set.
func (s *SubjectByTypeSet) Len() int {
	return len(s.byType)
}

// SubjectSetForType returns the subject set associated with the given subject type, if any.
func (s *SubjectByTypeSet) SubjectSetForType(rr *core.RelationReference) (SubjectSet, bool) {
	found, ok := s.byType[tuple.JoinRelRef(rr.Namespace, rr.Relation)]
	return found, ok
}

func wrapCaveat(caveat *core.ContextualizedCaveat) *core.CaveatExpression {
	if caveat == nil {
		return nil
	}

	return &core.CaveatExpression{
		OperationOrCaveat: &core.CaveatExpression_Caveat{
			Caveat: caveat,
		},
	}
}