github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datasets/subjectsetbyresourceid.go (about)

     1  package datasets
     2  
     3  import (
     4  	"fmt"
     5  
     6  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
     7  	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
     8  )
     9  
    10  // NewSubjectSetByResourceID creates and returns a map of subject sets, indexed by resource ID.
    11  func NewSubjectSetByResourceID() SubjectSetByResourceID {
    12  	return SubjectSetByResourceID{
    13  		subjectSetByResourceID: map[string]SubjectSet{},
    14  	}
    15  }
    16  
    17  // SubjectSetByResourceID defines a helper type which maps from a resource ID to its associated found
    18  // subjects, in the form of a subject set per resource ID.
    19  type SubjectSetByResourceID struct {
    20  	subjectSetByResourceID map[string]SubjectSet
    21  }
    22  
    23  func (ssr SubjectSetByResourceID) add(resourceID string, subject *v1.FoundSubject) error {
    24  	if subject == nil {
    25  		return fmt.Errorf("cannot add a nil subject to SubjectSetByResourceID")
    26  	}
    27  
    28  	_, ok := ssr.subjectSetByResourceID[resourceID]
    29  	if !ok {
    30  		ssr.subjectSetByResourceID[resourceID] = NewSubjectSet()
    31  	}
    32  	return ssr.subjectSetByResourceID[resourceID].Add(subject)
    33  }
    34  
    35  // ConcreteSubjectCount returns the number concrete subjects in the map.
    36  func (ssr SubjectSetByResourceID) ConcreteSubjectCount() int {
    37  	count := 0
    38  	for _, subjectSet := range ssr.subjectSetByResourceID {
    39  		count += subjectSet.ConcreteSubjectCount()
    40  	}
    41  	return count
    42  }
    43  
    44  // AddFromRelationship adds the subject found in the given relationship to this map, indexed at
    45  // the resource ID specified in the relationship.
    46  func (ssr SubjectSetByResourceID) AddFromRelationship(relationship *core.RelationTuple) error {
    47  	return ssr.add(relationship.ResourceAndRelation.ObjectId, &v1.FoundSubject{
    48  		SubjectId:        relationship.Subject.ObjectId,
    49  		CaveatExpression: wrapCaveat(relationship.Caveat),
    50  	})
    51  }
    52  
    53  // UnionWith unions the map's sets with the other map of sets provided.
    54  func (ssr SubjectSetByResourceID) UnionWith(other map[string]*v1.FoundSubjects) error {
    55  	for resourceID, subjects := range other {
    56  		if subjects == nil {
    57  			return fmt.Errorf("received nil FoundSubjects in other map of SubjectSetByResourceID's UnionWith for key %s", resourceID)
    58  		}
    59  
    60  		for _, subject := range subjects.FoundSubjects {
    61  			if err := ssr.add(resourceID, subject); err != nil {
    62  				return err
    63  			}
    64  		}
    65  	}
    66  
    67  	return nil
    68  }
    69  
    70  // IntersectionDifference performs an in-place intersection between the two maps' sets.
    71  func (ssr SubjectSetByResourceID) IntersectionDifference(other SubjectSetByResourceID) error {
    72  	for otherResourceID, otherSubjectSet := range other.subjectSetByResourceID {
    73  		existing, ok := ssr.subjectSetByResourceID[otherResourceID]
    74  		if !ok {
    75  			continue
    76  		}
    77  
    78  		err := existing.IntersectionDifference(otherSubjectSet)
    79  		if err != nil {
    80  			return err
    81  		}
    82  
    83  		if existing.IsEmpty() {
    84  			delete(ssr.subjectSetByResourceID, otherResourceID)
    85  		}
    86  	}
    87  
    88  	for existingResourceID := range ssr.subjectSetByResourceID {
    89  		_, ok := other.subjectSetByResourceID[existingResourceID]
    90  		if !ok {
    91  			delete(ssr.subjectSetByResourceID, existingResourceID)
    92  			continue
    93  		}
    94  	}
    95  
    96  	return nil
    97  }
    98  
    99  // SubtractAll subtracts all sets in the other map from this map's sets.
   100  func (ssr SubjectSetByResourceID) SubtractAll(other SubjectSetByResourceID) {
   101  	for otherResourceID, otherSubjectSet := range other.subjectSetByResourceID {
   102  		existing, ok := ssr.subjectSetByResourceID[otherResourceID]
   103  		if !ok {
   104  			continue
   105  		}
   106  
   107  		existing.SubtractAll(otherSubjectSet)
   108  		if existing.IsEmpty() {
   109  			delete(ssr.subjectSetByResourceID, otherResourceID)
   110  		}
   111  	}
   112  }
   113  
   114  // IsEmpty returns true if the map is empty.
   115  func (ssr SubjectSetByResourceID) IsEmpty() bool {
   116  	return len(ssr.subjectSetByResourceID) == 0
   117  }
   118  
   119  // AsMap converts the map into a map for storage in a proto.
   120  func (ssr SubjectSetByResourceID) AsMap() map[string]*v1.FoundSubjects {
   121  	mapped := make(map[string]*v1.FoundSubjects, len(ssr.subjectSetByResourceID))
   122  	for resourceID, subjectsSet := range ssr.subjectSetByResourceID {
   123  		mapped[resourceID] = subjectsSet.AsFoundSubjects()
   124  	}
   125  	return mapped
   126  }