Add support for basic cursors and limits to LookupSubjects

This change supports a limit (called the "concrete limit") on LookupSubjects and will filter concrete subjects based on the returned cursor. This change does *not* filter intermediate lookups, which will be done in a followup PR.
authzed · Jun 1, 2023 · da48a6b · da48a6b
1 parent fb656d6
commit da48a6b
Show file tree

Hide file tree

Showing 28 changed files with 2,936 additions and 627 deletions.
diff --git a/e2e/go.mod b/e2e/go.mod
@@ -3,7 +3,7 @@ module github.com/authzed/spicedb/e2e
 go 1.19
 
 require (
-	github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256
+	github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277
 	github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc
 	github.com/authzed/spicedb v1.21.0
 	github.com/brianvoe/gofakeit/v6 v6.15.0

diff --git a/e2e/go.sum b/e2e/go.sum
@@ -394,8 +394,8 @@ github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAE
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/antlr/antlr4/runtime/Go/antlr v1.4.10 h1:yL7+Jz0jTC6yykIK/Wh74gnTJnrGr5AyrNMXuA0gves=
 github.com/antlr/antlr4/runtime/Go/antlr v1.4.10/go.mod h1:F7bn7fEU90QkQ3tnmaTx3LTKLEDqnwWODIYppRQ5hnY=
-github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256 h1:fu0vsI/+dxpQpdWihrI219piaL2GbXOTrLAMM8ZjanM=
-github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
+github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277 h1:NKwiHNkS5UYymzADmrrJ/dfMRrN4kKpWpP9Fwx9Nids=
+github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
 github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc h1:FzQP1mljgX15vawHVHlPXQrB6F6wnR/6CftpED+2cvg=
 github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc/go.mod h1:erPFLN0tntt2V4PAxoUTwqtinm3UhFGzJrcxYKKzGUM=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=

diff --git a/go.mod b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/IBM/pgxpoolprometheus v1.1.1
 	github.com/Masterminds/squirrel v1.5.3
 	github.com/agnivade/wasmbrowsertest v0.7.0
-	github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256
+	github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277
 	github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc
 	github.com/aws/aws-sdk-go v1.44.110
 	github.com/benbjohnson/clock v1.3.3
@@ -63,6 +63,7 @@ require (
 	github.com/rs/zerolog v1.29.0
 	github.com/samber/lo v1.38.1
 	github.com/scylladb/go-set v1.0.2
+	github.com/sean-/sysexits v1.0.0
 	github.com/sercand/kuberesolver/v4 v4.0.0
 	github.com/shopspring/decimal v1.3.1
 	github.com/spf13/cobra v1.7.0
@@ -289,7 +290,6 @@ require (
 	github.com/sanposhiho/wastedassign/v2 v2.0.7 // indirect
 	github.com/sashamelentyev/interfacebloat v1.1.0 // indirect
 	github.com/sashamelentyev/usestdlibvars v1.23.0 // indirect
-	github.com/sean-/sysexits v1.0.0 // indirect
 	github.com/securego/gosec/v2 v2.15.0 // indirect
 	github.com/shabbyrobe/gocovmerge v0.0.0-20190829150210-3e036491d500 // indirect
 	github.com/shazow/go-diff v0.0.0-20160112020656-b6b7b6733b8c // indirect

diff --git a/go.sum b/go.sum
@@ -447,8 +447,8 @@ github.com/ashanbrown/forbidigo v1.5.1 h1:WXhzLjOlnuDYPYQo/eFlcFMi8X/kLfvWLYu6CS
 github.com/ashanbrown/forbidigo v1.5.1/go.mod h1:Y8j9jy9ZYAEHXdu723cUlraTqbzjKF1MUyfOKL+AjcU=
 github.com/ashanbrown/makezero v1.1.1 h1:iCQ87C0V0vSyO+M9E/FZYbu65auqH0lnsOkf5FcB28s=
 github.com/ashanbrown/makezero v1.1.1/go.mod h1:i1bJLCRSCHOcOa9Y6MyF2FTfMZMFdHvxKHxgO5Z1axI=
-github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256 h1:fu0vsI/+dxpQpdWihrI219piaL2GbXOTrLAMM8ZjanM=
-github.com/authzed/authzed-go v0.8.1-0.20230601141426-d411c3f66256/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
+github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277 h1:NKwiHNkS5UYymzADmrrJ/dfMRrN4kKpWpP9Fwx9Nids=
+github.com/authzed/authzed-go v0.8.1-0.20230601170210-d97cfb410277/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
 github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc h1:FzQP1mljgX15vawHVHlPXQrB6F6wnR/6CftpED+2cvg=
 github.com/authzed/grpcutil v0.0.0-20230509155820-7a6fedb71dbc/go.mod h1:erPFLN0tntt2V4PAxoUTwqtinm3UhFGzJrcxYKKzGUM=
 github.com/aws/aws-sdk-go v1.17.4/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=

diff --git a/internal/datasets/basesubjectset.go b/internal/datasets/basesubjectset.go
@@ -264,6 +264,25 @@ func (bss BaseSubjectSet[T]) AsSlice() []T {
 	return values
 }
 
+// SubjectCount returns the number of subjects in the set.
+func (bss BaseSubjectSet[T]) SubjectCount() int {
+	if _, ok := bss.wildcard.get(); ok {
+		return bss.ConcreteSubjectCount() + 1
+	}
+	return bss.ConcreteSubjectCount()
+}
+
+// ConcreteSubjectCount returns the number of concrete subjects in the set.
+func (bss BaseSubjectSet[T]) ConcreteSubjectCount() int {
+	return len(bss.concrete)
+}
+
+// HasWildcard returns true if the subject set contains the specialized wildcard subject.
+func (bss BaseSubjectSet[T]) HasWildcard() bool {
+	_, ok := bss.wildcard.get()
+	return ok
+}
+
 // Clone returns a clone of this subject set. Note that this is a shallow clone.
 // NOTE: Should only be used when performance is not a concern.
 func (bss BaseSubjectSet[T]) Clone() BaseSubjectSet[T] {

diff --git a/internal/datasets/subjectset_test.go b/internal/datasets/subjectset_test.go
@@ -243,6 +243,13 @@ func TestSubjectSetAdd(t *testing.T) {
 			expectedSet := tc.expectedSet
 			computedSet := existingSet.AsSlice()
 			testutil.RequireEquivalentSets(t, expectedSet, computedSet)
+
+			require.Equal(t, len(expectedSet), existingSet.SubjectCount())
+			if existingSet.HasWildcard() {
+				require.Equal(t, len(expectedSet), existingSet.ConcreteSubjectCount()+1)
+			} else {
+				require.Equal(t, len(expectedSet), existingSet.ConcreteSubjectCount())
+			}
 		})
 	}
 }

diff --git a/internal/datasets/subjectsetbyresourceid.go b/internal/datasets/subjectsetbyresourceid.go
@@ -32,6 +32,15 @@ func (ssr SubjectSetByResourceID) add(resourceID string, subject *v1.FoundSubjec
 	return ssr.subjectSetByResourceID[resourceID].Add(subject)
 }
 
+// ConcreteSubjectCount returns the number concrete subjects in the map.
+func (ssr SubjectSetByResourceID) ConcreteSubjectCount() int {
+	count := 0
+	for _, subjectSet := range ssr.subjectSetByResourceID {
+		count += subjectSet.ConcreteSubjectCount()
+	}
+	return count
+}
+
 // AddFromRelationship adds the subject found in the given relationship to this map, indexed at
 // the resource ID specified in the relationship.
 func (ssr SubjectSetByResourceID) AddFromRelationship(relationship *core.RelationTuple) error {

diff --git a/internal/datasets/subjectsetbyresourceid_test.go b/internal/datasets/subjectsetbyresourceid_test.go
@@ -43,6 +43,7 @@ func TestSubjectSetByResourceIDBasicOperations(t *testing.T) {
 	sort.Sort(sortFoundSubjects(asMap["seconddoc"].FoundSubjects))
 
 	require.Equal(t, expected, asMap)
+	require.Equal(t, 3, ssr.ConcreteSubjectCount())
 }
 
 func TestSubjectSetByResourceIDUnionWith(t *testing.T) {
@@ -88,6 +89,8 @@ func TestSubjectSetByResourceIDUnionWith(t *testing.T) {
 			},
 		},
 	}, found)
+
+	require.Equal(t, 5, ssr.ConcreteSubjectCount())
 }
 
 type sortFoundSubjects []*v1.FoundSubject

diff --git a/internal/datasets/subjectsetbytype.go b/internal/datasets/subjectsetbytype.go
@@ -56,6 +56,23 @@ func (s *SubjectByTypeSet) ForEachType(handler func(rr *core.RelationReference,
 	}
 }
 
+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) {