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.

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.20230620170737-8257e7bd388e
+	github.com/authzed/authzed-go v0.9.0
 	github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5
 	github.com/authzed/spicedb v1.21.0
 	github.com/brianvoe/gofakeit/v6 v6.15.0

e2e/go.sum

@@ -5,8 +5,8 @@ cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGB
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 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.20230620170737-8257e7bd388e h1:jVBWn6jQw2SnsVrlGcnKEq7cpqmj5lm1mMl1kzuaykQ=
-github.com/authzed/authzed-go v0.8.1-0.20230620170737-8257e7bd388e/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
+github.com/authzed/authzed-go v0.9.0 h1:FBWWwYiZrreGN94R9EEIy1S2s0UAH0Hn7MWBRtbtF+w=
+github.com/authzed/authzed-go v0.9.0/go.mod h1:9Pl5jDQJHrjbMDuCrsa+Q6Tqmi1f2pDdIn/qNGI++vA=
 github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5 h1:Fg92G8sNNODbNe2ckJoLeMEPeDqSfygmXnpEXDnVifU=
 github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5/go.mod h1:qx105brQubHFYLRja6wlHA+JB8DSK+yhb8uc8aFA5NQ=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/IBM/pgxpoolprometheus v1.1.1
 	github.com/Masterminds/squirrel v1.5.4
 	github.com/agnivade/wasmbrowsertest v0.7.0
-	github.com/authzed/authzed-go v0.8.1-0.20230620170737-8257e7bd388e
+	github.com/authzed/authzed-go v0.9.0
 	github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5
 	github.com/aws/aws-sdk-go v1.44.110
 	github.com/benbjohnson/clock v1.3.3

go.sum

@@ -452,8 +452,8 @@ github.com/ashanbrown/forbidigo v1.5.3 h1:jfg+fkm/snMx+V9FBwsl1d340BV/99kZGv5jN9
 github.com/ashanbrown/forbidigo v1.5.3/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.20230620170737-8257e7bd388e h1:jVBWn6jQw2SnsVrlGcnKEq7cpqmj5lm1mMl1kzuaykQ=
-github.com/authzed/authzed-go v0.8.1-0.20230620170737-8257e7bd388e/go.mod h1:qn4HCG0DQcLybaRePpVFW/Gvgz9UgkvobzyBoA5a49c=
+github.com/authzed/authzed-go v0.9.0 h1:FBWWwYiZrreGN94R9EEIy1S2s0UAH0Hn7MWBRtbtF+w=
+github.com/authzed/authzed-go v0.9.0/go.mod h1:9Pl5jDQJHrjbMDuCrsa+Q6Tqmi1f2pDdIn/qNGI++vA=
 github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5 h1:Fg92G8sNNODbNe2ckJoLeMEPeDqSfygmXnpEXDnVifU=
 github.com/authzed/grpcutil v0.0.0-20230703173955-bdd0ac3f16a5/go.mod h1:qx105brQubHFYLRja6wlHA+JB8DSK+yhb8uc8aFA5NQ=
 github.com/aws/aws-sdk-go v1.17.4/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=

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] {

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())
+			}
 		})
 	}
 }

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 {

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

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) {

internal/dispatch/graph/lookupresources_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/authzed/spicedb/internal/dispatch"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	"github.com/authzed/spicedb/internal/testfixtures"
+	"github.com/authzed/spicedb/internal/testutil"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
@@ -351,53 +352,6 @@ func (a OrderedResolved) Less(i, j int) bool {
 
 func (a OrderedResolved) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 
-func joinTuples(first []*core.RelationTuple, second []*core.RelationTuple) []*core.RelationTuple {
-	return append(first, second...)
-}
-
-func genTuplesWithOffset(resourceName string, relation string, subjectName string, subjectID string, offset int, number int) []*core.RelationTuple {
-	return genTuplesWithCaveat(resourceName, relation, subjectName, subjectID, "", nil, offset, number)
-}
-
-func genTuples(resourceName string, relation string, subjectName string, subjectID string, number int) []*core.RelationTuple {
-	return genTuplesWithOffset(resourceName, relation, subjectName, subjectID, 0, number)
-}
-
-func genSubjectTuples(resourceName string, relation string, subjectName string, subjectRelation string, number int) []*core.RelationTuple {
-	tuples := make([]*core.RelationTuple, 0, number)
-	for i := 0; i < number; i++ {
-		tpl := &core.RelationTuple{
-			ResourceAndRelation: ONR(resourceName, fmt.Sprintf("%s-%d", resourceName, i), relation),
-			Subject:             ONR(subjectName, fmt.Sprintf("%s-%d", subjectName, i), subjectRelation),
-		}
-		tuples = append(tuples, tpl)
-	}
-	return tuples
-}
-
-func genTuplesWithCaveat(resourceName string, relation string, subjectName string, subjectID string, caveatName string, context map[string]any, offset int, number int) []*core.RelationTuple {
-	tuples := make([]*core.RelationTuple, 0, number)
-	for i := 0; i < number; i++ {
-		tpl := &core.RelationTuple{
-			ResourceAndRelation: ONR(resourceName, fmt.Sprintf("%s-%d", resourceName, i+offset), relation),
-			Subject:             ONR(subjectName, subjectID, "..."),
-		}
-		if caveatName != "" {
-			tpl = tuple.MustWithCaveat(tpl, caveatName, context)
-		}
-		tuples = append(tuples, tpl)
-	}
-	return tuples
-}
-
-func genResourceIds(resourceName string, number int) []string {
-	resourceIDs := make([]string, 0, number)
-	for i := 0; i < number; i++ {
-		resourceIDs = append(resourceIDs, fmt.Sprintf("%s-%d", resourceName, i))
-	}
-	return resourceIDs
-}
-
 func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 	testCases := []struct {
 		name                string
@@ -416,13 +370,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user
 				permission view = viewer + editor
   			 }`,
-			joinTuples(
-				genTuples("document", "viewer", "user", "tom", 1510),
-				genTuples("document", "editor", "user", "tom", 1510),
+			testutil.JoinTuples(
+				testutil.GenTuples("document", "viewer", "user", "tom", 1510),
+				testutil.GenTuples("document", "editor", "user", "tom", 1510),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 1510),
+			testutil.GenResourceIds("document", 1510),
 		},
 		{
 			"basic exclusion",
@@ -433,10 +387,10 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user
 				permission view = viewer - banned
   			 }`,
-			genTuples("document", "viewer", "user", "tom", 1010),
+			testutil.GenTuples("document", "viewer", "user", "tom", 1010),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 1010),
+			testutil.GenResourceIds("document", 1010),
 		},
 		{
 			"basic intersection",
@@ -447,13 +401,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user
 				permission view = viewer & editor
   			 }`,
-			joinTuples(
-				genTuples("document", "viewer", "user", "tom", 510),
-				genTuples("document", "editor", "user", "tom", 510),
+			testutil.JoinTuples(
+				testutil.GenTuples("document", "viewer", "user", "tom", 510),
+				testutil.GenTuples("document", "editor", "user", "tom", 510),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 510),
+			testutil.GenResourceIds("document", 510),
 		},
 		{
 			"union and exclused union",
@@ -466,13 +420,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				permission can_view = viewer - banned
 				permission view = can_view + editor
   			 }`,
-			joinTuples(
-				genTuples("document", "viewer", "user", "tom", 1310),
-				genTuplesWithOffset("document", "editor", "user", "tom", 1250, 1200),
+			testutil.JoinTuples(
+				testutil.GenTuples("document", "viewer", "user", "tom", 1310),
+				testutil.GenTuplesWithOffset("document", "editor", "user", "tom", 1250, 1200),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 2450),
+			testutil.GenResourceIds("document", 2450),
 		},
 		{
 			"basic caveats",
@@ -486,10 +440,10 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user with somecaveat
 				permission view = viewer
   			 }`,
-			genTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{"somecondition": 42}, 0, 2450),
+			testutil.GenTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{"somecondition": 42}, 0, 2450),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 2450),
+			testutil.GenResourceIds("document", 2450),
 		},
 		{
 			"excluded items",
@@ -500,13 +454,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user
 				permission view = viewer - banned
   			 }`,
-			joinTuples(
-				genTuples("document", "viewer", "user", "tom", 1310),
-				genTuplesWithOffset("document", "banned", "user", "tom", 1210, 100),
+			testutil.JoinTuples(
+				testutil.GenTuples("document", "viewer", "user", "tom", 1310),
+				testutil.GenTuplesWithOffset("document", "banned", "user", "tom", 1210, 100),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 1210),
+			testutil.GenResourceIds("document", 1210),
 		},
 		{
 			"basic caveats with missing field",
@@ -520,10 +474,10 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user with somecaveat
 				permission view = viewer
   			 }`,
-			genTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{}, 0, 2450),
+			testutil.GenTuplesWithCaveat("document", "viewer", "user", "tom", "somecaveat", map[string]any{}, 0, 2450),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 2450),
+			testutil.GenResourceIds("document", 2450),
 		},
 		{
 			"larger arrow dispatch",
@@ -537,13 +491,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation folder: folder
 				permission view = folder->viewer
   			 }`,
-			joinTuples(
-				genTuples("folder", "viewer", "user", "tom", 150),
-				genSubjectTuples("document", "folder", "folder", "...", 150),
+			testutil.JoinTuples(
+				testutil.GenTuples("folder", "viewer", "user", "tom", 150),
+				testutil.GenSubjectTuples("document", "folder", "folder", "...", 150),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 150),
+			testutil.GenResourceIds("document", 150),
 		},
 		{
 			"big",
@@ -554,13 +508,13 @@ func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 				relation viewer: user
 				permission view = viewer + editor
   			 }`,
-			joinTuples(
-				genTuples("document", "viewer", "user", "tom", 15100),
-				genTuples("document", "editor", "user", "tom", 15100),
+			testutil.JoinTuples(
+				testutil.GenTuples("document", "viewer", "user", "tom", 15100),
+				testutil.GenTuples("document", "editor", "user", "tom", 15100),
 			),
 			RR("document", "view"),
 			ONR("user", "tom", "..."),
-			genResourceIds("document", 15100),
+			testutil.GenResourceIds("document", 15100),
 		},
 	}