Merge commit from fork

* fix: fix entrypoints over relations that are reused for arrows

Fixes the reachability graph to prevent skipping of necessary entrypoint computations when a relation is reused for an arrow in the same permission

* fix: Fix building query planner trees that contain arrows that may
contain types that don't have the right side of the arrow.

---------

Co-authored-by: Barak Michener <[email protected]>

Author: josephschorr

internal/services/integrationtesting/testconfigs/arrowsublr.yaml

@@ -0,0 +1,30 @@
+schema: |+
+  definition special_user {}
+
+  definition user {
+  	relation special_user_mapping: special_user
+  	permission special_user = special_user_mapping
+  }
+
+  definition group {
+  	relation member: user
+  	permission membership = member + member->special_user
+  }
+
+  definition system {
+  	relation viewer: user | group#membership
+  	permission view = viewer + viewer->special_user
+  }
+
+
+relationships: |-
+  system:somesystem#viewer@group:somegroup#membership
+
+  group:somegroup#member@user:someuser1
+
+  user:someuser1#special_user_mapping@special_user:specialuser
+assertions:
+  assertTrue:
+    - system:somesystem#view@special_user:specialuser
+  assertFalse: []
+validation: {}

pkg/query/build_tree.go

@@ -1,6 +1,7 @@
 package query
 
 import (
+	"errors"
 	"fmt"
 
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -89,7 +90,10 @@ func (b *iteratorBuilder) buildIteratorFromSchemaInternal(definitionName string,
 	} else if r, ok := def.GetRelation(relationName); ok {
 		result, err = b.buildIteratorFromRelation(r, withSubRelations)
 	} else {
-		err = fmt.Errorf("BuildIteratorFromSchema: couldn't find a relation or permission named `%s` in definition `%s`", relationName, definitionName)
+		err = RelationNotFoundError{
+			definitionName: definitionName,
+			relationName:   relationName,
+		}
 	}
 
 	// Remove from building after we're done (allows reuse in other branches)
@@ -284,36 +288,76 @@ func (b *iteratorBuilder) buildBaseRelationIterator(br *schema.BaseRelation, wit
 // buildArrowIterators creates a union of arrow iterators for the given relation and right-hand side
 func (b *iteratorBuilder) buildArrowIterators(rel *schema.Relation, rightSide string) (Iterator, error) {
 	union := NewUnion()
+	hasMultipleBaseRelations := len(rel.BaseRelations()) > 1
+	var lastNotFoundError error
+
 	for _, br := range rel.BaseRelations() {
 		left, err := b.buildBaseRelationIterator(br, false)
 		if err != nil {
 			return nil, err
 		}
 		right, err := b.buildIteratorFromSchemaInternal(br.Type(), rightSide, false)
 		if err != nil {
+			// If the right side doesn't exist on this type, the arrow produces an empty set.
+			// This is valid when a relation has multiple types and the arrow only
+			// applies to some of them. If there's only one base relation, we should error.
+			if errors.As(err, &RelationNotFoundError{}) {
+				if hasMultipleBaseRelations {
+					union.addSubIterator(NewEmptyFixedIterator())
+					continue
+				}
+				lastNotFoundError = err
+				continue
+			}
 			return nil, err
 		}
 		arrow := NewArrow(left, right)
 		union.addSubIterator(arrow)
 	}
+
+	// If we have no sub-iterators and only have a not-found error, return that error
+	if len(union.Subiterators()) == 0 && lastNotFoundError != nil {
+		return nil, lastNotFoundError
+	}
+
 	return union, nil
 }
 
 // buildIntersectionArrowIterators creates a union of intersection arrow iterators for the given relation and right-hand side
 func (b *iteratorBuilder) buildIntersectionArrowIterators(rel *schema.Relation, rightSide string) (Iterator, error) {
 	union := NewUnion()
+	hasMultipleBaseRelations := len(rel.BaseRelations()) > 1
+	var lastNotFoundError error
+
 	for _, br := range rel.BaseRelations() {
 		left, err := b.buildBaseRelationIterator(br, false)
 		if err != nil {
 			return nil, err
 		}
 		right, err := b.buildIteratorFromSchemaInternal(br.Type(), rightSide, false)
 		if err != nil {
+			// If the right side doesn't exist on this type, the intersection arrow produces an empty set.
+			// This is valid when a relation has multiple types and the arrow only
+			// applies to some of them. If there's only one base relation, we should error.
+			if errors.As(err, &RelationNotFoundError{}) {
+				if hasMultipleBaseRelations {
+					union.addSubIterator(NewEmptyFixedIterator())
+					continue
+				}
+				lastNotFoundError = err
+				continue
+			}
 			return nil, err
 		}
 		intersectionArrow := NewIntersectionArrow(left, right)
 		union.addSubIterator(intersectionArrow)
 	}
+
+	// If we have no sub-iterators and only have a not-found error, return that error
+	if len(union.Subiterators()) == 0 && lastNotFoundError != nil {
+		return nil, lastNotFoundError
+	}
+
 	return union, nil
 }
 
@@ -327,3 +371,13 @@ func functionTypeString(ft schema.FunctionType) string {
 		return "unknown"
 	}
 }
+
+// RelationNotFoundError is returned when a relation or permission is not found in a definition
+type RelationNotFoundError struct {
+	definitionName string
+	relationName   string
+}
+
+func (e RelationNotFoundError) Error() string {
+	return fmt.Sprintf("BuildIteratorFromSchema: couldn't find a relation or permission named `%s` in definition `%s`", e.relationName, e.definitionName)
+}

pkg/schema/reachabilitygraph.go

@@ -37,7 +37,7 @@ func (rg *DefinitionReachability) RelationsEncounteredForResource(
 	ctx context.Context,
 	resourceType *core.RelationReference,
 ) ([]*core.RelationReference, error) {
-	_, relationNames, err := rg.computeEntrypoints(ctx, resourceType, nil /* include all entrypoints */, reachabilityFull, entrypointLookupFindAll)
+	_, relationNames, err := rg.computeEntrypoints(ctx, resourceType, nil, reachabilityFull /* include all entrypoints */, entrypointLookupFindAll)
 	if err != nil {
 		return nil, err
 	}
@@ -85,11 +85,12 @@ func (rg *DefinitionReachability) RelationsEncounteredForSubject(
 						continue
 					}
 
-					encounteredRelations := map[string]struct{}{}
+					allEncounteredRelations := mapz.NewSet[string]()
+					encounteredRelationsForComputation := mapz.NewSet[string]()
 					err := nrg.collectEntrypoints(ctx, &core.RelationReference{
 						Namespace: nsDef.Name,
 						Relation:  relation.Name,
-					}, subjectType, collected, encounteredRelations, reachabilityFull, entrypointLookupFindAll)
+					}, subjectType, collected, allEncounteredRelations, encounteredRelationsForComputation, reachabilityFull, entrypointLookupFindAll)
 					if err != nil {
 						return nil, err
 					}
@@ -193,10 +194,13 @@ func (rg *DefinitionReachability) computeEntrypoints(
 	}
 
 	collected := &[]ReachabilityEntrypoint{}
-	encounteredRelations := map[string]struct{}{}
-	err := rg.collectEntrypoints(ctx, resourceType, optionalSubjectType, collected, encounteredRelations, reachabilityOption, entrypointLookupOption)
+
+	allEncounteredRelations := mapz.NewSet[string]()
+	encounteredRelationsForComputation := mapz.NewSet[string]()
+
+	err := rg.collectEntrypoints(ctx, resourceType, optionalSubjectType, collected, allEncounteredRelations, encounteredRelationsForComputation, reachabilityOption, entrypointLookupOption)
 	if err != nil {
-		return nil, slices.Collect(maps.Keys(encounteredRelations)), err
+		return nil, allEncounteredRelations.AsSlice(), err
 	}
 
 	collectedEntrypoints := *collected
@@ -213,7 +217,7 @@ func (rg *DefinitionReachability) computeEntrypoints(
 	for _, entrypoint := range collectedEntrypoints {
 		hash, err := entrypoint.Hash()
 		if err != nil {
-			return nil, slices.Collect(maps.Keys(encounteredRelations)), err
+			return nil, allEncounteredRelations.AsSlice(), err
 		}
 
 		if _, ok := entrypointMap[hash]; !ok {
@@ -222,7 +226,7 @@ func (rg *DefinitionReachability) computeEntrypoints(
 		}
 	}
 
-	return uniqueEntrypoints, slices.Collect(maps.Keys(encounteredRelations)), nil
+	return uniqueEntrypoints, allEncounteredRelations.AsSlice(), nil
 }
 
 func (rg *DefinitionReachability) getOrBuildGraph(ctx context.Context, resourceType *core.RelationReference, reachabilityOption reachabilityOption) (*core.ReachabilityGraph, error) {
@@ -253,17 +257,17 @@ func (rg *DefinitionReachability) collectEntrypoints(
 	resourceType *core.RelationReference,
 	optionalSubjectType *core.RelationReference,
 	collected *[]ReachabilityEntrypoint,
-	encounteredRelations map[string]struct{},
+	allEncounteredRelations *mapz.Set[string],
+	encounteredRelationsForComputation *mapz.Set[string],
 	reachabilityOption reachabilityOption,
 	entrypointLookupOption entrypointLookupOption,
 ) error {
 	// Ensure that we only process each relation once.
 	key := tuple.JoinRelRef(resourceType.Namespace, resourceType.Relation)
-	if _, ok := encounteredRelations[key]; ok {
+	if !encounteredRelationsForComputation.Add(key) {
 		return nil
 	}
-
-	encounteredRelations[key] = struct{}{}
+	allEncounteredRelations.Add(key)
 
 	rrg, err := rg.getOrBuildGraph(ctx, resourceType, reachabilityOption)
 	if err != nil {
@@ -274,7 +278,7 @@ func (rg *DefinitionReachability) collectEntrypoints(
 		// Add subject type entrypoints.
 		subjectTypeEntrypoints, ok := rrg.EntrypointsBySubjectType[optionalSubjectType.Namespace]
 		if ok {
-			addEntrypoints(subjectTypeEntrypoints, resourceType, collected, encounteredRelations)
+			addEntrypoints(subjectTypeEntrypoints, resourceType, collected, allEncounteredRelations, encounteredRelationsForComputation)
 		}
 
 		if entrypointLookupOption == entrypointLookupFindOne && len(*collected) > 0 {
@@ -284,7 +288,7 @@ func (rg *DefinitionReachability) collectEntrypoints(
 		// Add subject relation entrypoints.
 		subjectRelationEntrypoints, ok := rrg.EntrypointsBySubjectRelation[tuple.JoinRelRef(optionalSubjectType.Namespace, optionalSubjectType.Relation)]
 		if ok {
-			addEntrypoints(subjectRelationEntrypoints, resourceType, collected, encounteredRelations)
+			addEntrypoints(subjectRelationEntrypoints, resourceType, collected, allEncounteredRelations, encounteredRelationsForComputation)
 		}
 
 		if entrypointLookupOption == entrypointLookupFindOne && len(*collected) > 0 {
@@ -293,11 +297,11 @@ func (rg *DefinitionReachability) collectEntrypoints(
 	} else {
 		// Add all entrypoints.
 		for _, entrypoints := range rrg.EntrypointsBySubjectType {
-			addEntrypoints(entrypoints, resourceType, collected, encounteredRelations)
+			addEntrypoints(entrypoints, resourceType, collected, allEncounteredRelations, encounteredRelationsForComputation)
 		}
 
 		for _, entrypoints := range rrg.EntrypointsBySubjectRelation {
-			addEntrypoints(entrypoints, resourceType, collected, encounteredRelations)
+			addEntrypoints(entrypoints, resourceType, collected, allEncounteredRelations, encounteredRelationsForComputation)
 		}
 	}
 
@@ -309,7 +313,7 @@ func (rg *DefinitionReachability) collectEntrypoints(
 	for _, entrypointSetKey := range keys {
 		entrypointSet := rrg.EntrypointsBySubjectRelation[entrypointSetKey]
 		if entrypointSet.SubjectRelation != nil && entrypointSet.SubjectRelation.Relation != tuple.Ellipsis {
-			err := rg.collectEntrypoints(ctx, entrypointSet.SubjectRelation, optionalSubjectType, collected, encounteredRelations, reachabilityOption, entrypointLookupOption)
+			err := rg.collectEntrypoints(ctx, entrypointSet.SubjectRelation, optionalSubjectType, collected, allEncounteredRelations, encounteredRelationsForComputation, reachabilityOption, entrypointLookupOption)
 			if err != nil {
 				return err
 			}
@@ -323,13 +327,12 @@ func (rg *DefinitionReachability) collectEntrypoints(
 	return nil
 }
 
-func addEntrypoints(entrypoints *core.ReachabilityEntrypoints, parentRelation *core.RelationReference, collected *[]ReachabilityEntrypoint, encounteredRelations map[string]struct{}) {
+func addEntrypoints(entrypoints *core.ReachabilityEntrypoints, parentRelation *core.RelationReference, collected *[]ReachabilityEntrypoint, allEncounteredRelations *mapz.Set[string], encounteredRelationsForComputation *mapz.Set[string]) {
 	for _, entrypoint := range entrypoints.Entrypoints {
 		if entrypoint.TuplesetRelation != "" {
 			key := tuple.JoinRelRef(entrypoint.TargetRelation.Namespace, entrypoint.TuplesetRelation)
-			encounteredRelations[key] = struct{}{}
+			allEncounteredRelations.Add(key)
 		}
-
 		*collected = append(*collected, ReachabilityEntrypoint{entrypoint, parentRelation})
 	}
 }

pkg/schema/reachabilitygraph_test.go

@@ -1180,6 +1180,62 @@ func TestReachabilityGraph(t *testing.T) {
 				rrt("document", "view", false),
 			},
 		},
+		{
+			"multiple arrows with same starting point",
+			`definition special_user {}
+
+  definition user {
+  	relation special_user_mapping: special_user
+  	permission special_user = special_user_mapping
+  }
+
+  definition group {
+  	relation member: user
+  	permission membership = member + member->special_user
+  }
+
+  definition system {
+  	relation viewer: user | group#membership
+  	permission view = viewer + viewer->special_user
+  }`,
+			rr("system", "view"),
+			rr("special_user", "..."),
+			[]rrtStruct{
+				rrt("user", "special_user_mapping", true),
+			},
+			[]rrtStruct{
+				rrt("user", "special_user_mapping", true),
+			},
+		},
+		{
+			"multiple arrows with same non-terminal starting point",
+			`definition special_user {}
+
+  definition user {
+  	relation special_user_mapping: special_user
+  	permission special_user = special_user_mapping
+  }
+
+  definition group {
+  	relation member: user
+  	permission membership = member + member->special_user
+  }
+
+  definition system {
+  	relation viewer: user | group#membership
+  	permission view = viewer + viewer->special_user
+  }`,
+			rr("system", "view"),
+			rr("user", "special_user"),
+			[]rrtStruct{
+				rrt("system", "view", true),
+				rrt("group", "membership", true),
+			},
+			[]rrtStruct{
+				rrt("system", "view", true),
+				rrt("group", "membership", true),
+			},
+		},
 	}
 
 	for _, tc := range testCases {