fix: improve LR consistency and support multiple resourcetypes (#2875)

Author: barakmich

pkg/query/alias.go

@@ -180,9 +180,32 @@ func (a *Alias) IterResourcesImpl(ctx *Context, subject ObjectAndRelation, filte
 		return nil, err
 	}
 
-	// If the relation on the alias iterator is the same as the subject relation,
-	// we have a self-edge and should yield it
-	shouldAddSelfEdge := a.relation == subject.Relation
+	// Check if we should add a self-edge (identity check: permission always grants access to itself)
+	// This matches LookupResources3's logic where subject type+relation must match resource type+relation
+	shouldAddSelfEdge := false
+	if a.relation == subject.Relation {
+		// Get the resource types from the iterator
+		resourceTypes, err := a.ResourceType()
+		if err != nil {
+			return nil, err
+		}
+
+		// Add self-edge if:
+		// - No resource types defined (empty/unconstrained iterator), OR
+		// - Subject type matches one of the possible resource types
+		// This allows self-edges for empty iterators while preventing them in nested contexts
+		// where the types don't match
+		if len(resourceTypes) == 0 {
+			shouldAddSelfEdge = true
+		} else {
+			for _, rt := range resourceTypes {
+				if rt.Type == subject.ObjectType {
+					shouldAddSelfEdge = true
+					break
+				}
+			}
+		}
+	}
 
 	return a.maybePrependSelfEdge(GetObject(subject), subSeq, shouldAddSelfEdge), nil
 }
@@ -215,7 +238,7 @@ func (a *Alias) ID() string {
 	return a.id
 }
 
-func (a *Alias) ResourceType() (ObjectType, error) {
+func (a *Alias) ResourceType() ([]ObjectType, error) {
 	return a.subIt.ResourceType()
 }
 

pkg/query/alias_test.go

@@ -760,7 +760,8 @@ func TestAlias_Types(t *testing.T) {
 
 		resourceType, err := alias.ResourceType()
 		require.NoError(err)
-		require.Equal("document", resourceType.Type)
+		require.Len(resourceType, 1)
+		require.Equal("document", resourceType[0].Type)
 	})
 
 	t.Run("SubjectTypes", func(t *testing.T) {

pkg/query/arrow.go

@@ -282,6 +282,11 @@ func (a *Arrow) IterResourcesImpl(ctx *Context, subject ObjectAndRelation, filte
 				yield(Path{}, err)
 				return
 			}
+
+			// Note: We used to filter self-edges here, but self-edges from Alias represent valid identity checks
+			// (e.g., team:first#member accessing team:first via member). Removing the filter allows these
+			// identity relationships to propagate through arrows correctly.
+
 			rightPathCount++
 
 			// For each right resource, get resources from left side
@@ -359,7 +364,7 @@ func (a *Arrow) ID() string {
 	return a.id
 }
 
-func (a *Arrow) ResourceType() (ObjectType, error) {
+func (a *Arrow) ResourceType() ([]ObjectType, error) {
 	// Arrow's resources come from the left side
 	return a.left.ResourceType()
 }

pkg/query/arrow_test.go

@@ -713,7 +713,8 @@ func TestArrow_Types(t *testing.T) {
 
 		resourceType, err := arrow.ResourceType()
 		require.NoError(err)
-		require.Equal("document", resourceType.Type) // From left iterator
+		require.Len(resourceType, 1)
+		require.Equal("document", resourceType[0].Type) // From left iterator
 	})
 
 	t.Run("SubjectTypes", func(t *testing.T) {

pkg/query/caveat.go

@@ -208,7 +208,7 @@ func (c *CaveatIterator) ID() string {
 	return c.id
 }
 
-func (c *CaveatIterator) ResourceType() (ObjectType, error) {
+func (c *CaveatIterator) ResourceType() ([]ObjectType, error) {
 	// Delegate to the wrapped iterator
 	return c.subiterator.ResourceType()
 }

pkg/query/caveat_test.go

@@ -550,7 +550,8 @@ func TestCaveatIterator_Types(t *testing.T) {
 
 		resourceType, err := caveatIter.ResourceType()
 		require.NoError(err)
-		require.Equal("document", resourceType.Type) // From subiterator
+		require.Len(resourceType, 1)
+		require.Equal("document", resourceType[0].Type) // From subiterator
 	})
 
 	t.Run("SubjectTypes", func(t *testing.T) {

pkg/query/datastore.go

@@ -270,9 +270,18 @@ func (r *RelationIterator) IterResourcesImpl(ctx *Context, subject ObjectAndRela
 		return EmptyPathSeq(), nil
 	}
 
+	// Handle wildcards first - they don't have subrelations and match any query relation
 	if r.base.Wildcard() {
 		return r.iterResourcesWildcardImpl(ctx, subject)
 	}
+
+	// Check if subject relation matches what this iterator expects.
+	// Both the schema's expected subrelation and the query's subject relation must match exactly.
+	// Ellipsis is a specific relation value, not a wildcard.
+	if r.base.Subrelation() != subject.Relation {
+		return EmptyPathSeq(), nil
+	}
+
 	filter := datastore.RelationshipsFilter{
 		OptionalResourceType:     r.base.DefinitionName(),
 		OptionalResourceRelation: r.base.RelationName(),
@@ -354,11 +363,11 @@ func (r *RelationIterator) ID() string {
 	return r.id
 }
 
-func (r *RelationIterator) ResourceType() (ObjectType, error) {
-	return ObjectType{
+func (r *RelationIterator) ResourceType() ([]ObjectType, error) {
+	return []ObjectType{{
 		Type:        r.base.DefinitionName(),
-		Subrelation: r.base.RelationName(),
-	}, nil
+		Subrelation: tuple.Ellipsis,
+	}}, nil
 }
 
 func (r *RelationIterator) SubjectTypes() ([]ObjectType, error) {

pkg/query/datastore_test.go

@@ -358,8 +358,9 @@ func TestRelationIterator_Types(t *testing.T) {
 
 		resourceType, err := relationIter.ResourceType()
 		require.NoError(err)
-		require.Equal("document", resourceType.Type)
-		require.Equal("viewer", resourceType.Subrelation)
+		require.Len(resourceType, 1)
+		require.Equal("document", resourceType[0].Type)
+		require.Equal(tuple.Ellipsis, resourceType[0].Subrelation)
 	})
 
 	t.Run("SubjectTypes_NoSubrelation", func(t *testing.T) {

pkg/query/exclusion.go

@@ -302,7 +302,7 @@ func (e *Exclusion) ID() string {
 	return e.id
 }
 
-func (e *Exclusion) ResourceType() (ObjectType, error) {
+func (e *Exclusion) ResourceType() ([]ObjectType, error) {
 	// Exclusion's resources come from the main set
 	return e.mainSet.ResourceType()
 }

pkg/query/exclusion_test.go

@@ -873,7 +873,8 @@ func TestExclusion_Types(t *testing.T) {
 
 		resourceType, err := exclusion.ResourceType()
 		require.NoError(err)
-		require.Equal("document", resourceType.Type) // From mainSet
+		require.Len(resourceType, 1)
+		require.Equal("document", resourceType[0].Type) // From mainSet
 	})
 
 	t.Run("SubjectTypes", func(t *testing.T) {