Merge pull request #2565 from authzed/barakmich/caveat_iterator

feat: Handle caveats and intersection arrows in the query plan

Author: barakmich

internal/datastore/memdb/readonly.go

@@ -493,11 +493,11 @@ func newSubjectSortedIterator(now time.Time, it memdb.ResultIterator, limit *uin
 		}
 
 		if skipCaveats && rt.OptionalCaveat != nil {
-			return nil, spiceerrors.MustBugf("unexpected caveat in result for relationship: %v", rt)
+			continue
 		}
 
 		if skipExpiration && rt.OptionalExpiration != nil {
-			return nil, spiceerrors.MustBugf("unexpected expiration in result for relationship: %v", rt)
+			continue
 		}
 
 		results = append(results, rt)
@@ -558,13 +558,11 @@ func newMemdbTupleIterator(now time.Time, it memdb.ResultIterator, limit *uint64
 			}
 
 			if skipCaveats && rt.OptionalCaveat != nil {
-				yield(rt, fmt.Errorf("unexpected caveat in result for relationship: %v", rt))
-				return
+				continue
 			}
 
 			if skipExpiration && rt.OptionalExpiration != nil {
-				yield(rt, fmt.Errorf("unexpected expiration in result for relationship: %v", rt))
-				return
+				continue
 			}
 
 			if rt.OptionalExpiration != nil && rt.OptionalExpiration.Before(now) {

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -4,14 +4,17 @@
 package integrationtesting_test
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
+	"github.com/authzed/spicedb/internal/caveats"
 	"github.com/authzed/spicedb/internal/datastore/dsfortesting"
 	"github.com/authzed/spicedb/internal/datastore/memdb"
 	"github.com/authzed/spicedb/internal/services/integrationtesting/consistencytestutil"
@@ -50,10 +53,12 @@ type queryPlanConsistencyHandle struct {
 
 func (q *queryPlanConsistencyHandle) buildContext(t *testing.T) *query.Context {
 	return &query.Context{
-		Context:   t.Context(),
-		Executor:  query.LocalExecutor{},
-		Datastore: q.ds,
-		Revision:  q.revision,
+		Context:      t.Context(),
+		Executor:     query.LocalExecutor{},
+		Datastore:    q.ds,
+		Revision:     q.revision,
+		CaveatRunner: caveats.NewCaveatRunner(caveattypes.Default.TypeSet),
+		TraceLogger:  query.NewTraceLogger(), // Enable tracing for debugging
 	}
 }
 
@@ -115,16 +120,32 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 							it, err := query.BuildIteratorFromSchema(handle.schema, rel.Resource.ObjectType, rel.Resource.Relation)
 							require.NoError(err)
 
-							t.Log(it.Explain())
-
 							qctx := handle.buildContext(t)
 
+							// Add caveat context from assertion if available
+							if len(assertion.CaveatContext) > 0 {
+								qctx.CaveatContext = assertion.CaveatContext
+							}
+
 							seq, err := qctx.Check(it, []query.Object{query.GetObject(rel.Resource)}, rel.Subject)
 							require.NoError(err)
 
 							rels, err := query.CollectAll(seq)
 							require.NoError(err)
 
+							// Print trace if test fails
+							if qctx.TraceLogger != nil {
+								defer func() {
+									if t.Failed() {
+										t.Logf("Trace for %s:\n%s", entry.name, qctx.TraceLogger.DumpTrace())
+										// Also print the tree structure for debugging
+										if it != nil {
+											t.Logf("Tree structure:\n%s", explainTree(it, 0))
+										}
+									}
+								}()
+							}
+
 							switch entry.expectedPermissionship {
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_CONDITIONAL_PERMISSION:
 								require.Len(rels, 1)
@@ -133,6 +154,9 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 								require.Len(rels, 1)
 								require.Nil(rels[0].Caveat)
 							case v1.CheckPermissionResponse_PERMISSIONSHIP_NO_PERMISSION:
+								if len(rels) != 0 && qctx.TraceLogger != nil {
+									t.Logf("Expected 0 relations but got %d. Trace:\n%s", len(rels), qctx.TraceLogger.DumpTrace())
+								}
 								require.Len(rels, 0)
 							}
 						})
@@ -142,3 +166,21 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 		}
 	})
 }
+
+// explainTree recursively explains the tree structure for debugging
+func explainTree(iter query.Iterator, depth int) string {
+	indent := strings.Repeat("  ", depth)
+	explain := iter.Explain()
+	result := fmt.Sprintf("%s%s: %s\n", indent, explain.Name, explain.Info)
+
+	for _, subExplain := range explain.SubExplain {
+		// For SubExplain, we need to create a dummy iterator to get the tree structure
+		// This is a simplified approach - in practice we'd need access to the actual sub-iterators
+		subResult := fmt.Sprintf("%s  %s: %s\n", indent, subExplain.Name, subExplain.Info)
+		result += subResult
+		// Note: We can't recursively call explainTree on SubExplain because it's not an Iterator
+		// This gives us one level of detail which should be sufficient for debugging
+	}
+
+	return result
+}

pkg/query/alias.go

@@ -34,7 +34,7 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 			}
 
 			// Also get relations from sub-iterator
-			subSeq, err := a.subIt.CheckImpl(ctx, resources, subject)
+			subSeq, err := ctx.Check(a.subIt, resources, subject)
 			if err != nil {
 				return nil, err
 			}
@@ -62,7 +62,7 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 	}
 
 	// No self-edge detected, just rewrite paths from sub-iterator
-	subSeq, err := a.subIt.CheckImpl(ctx, resources, subject)
+	subSeq, err := ctx.Check(a.subIt, resources, subject)
 	if err != nil {
 		return nil, err
 	}
@@ -83,7 +83,7 @@ func (a *Alias) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 }
 
 func (a *Alias) IterSubjectsImpl(ctx *Context, resource Object) (PathSeq, error) {
-	subSeq, err := a.subIt.IterSubjectsImpl(ctx, resource)
+	subSeq, err := ctx.IterSubjects(a.subIt, resource)
 	if err != nil {
 		return nil, err
 	}
@@ -104,7 +104,7 @@ func (a *Alias) IterSubjectsImpl(ctx *Context, resource Object) (PathSeq, error)
 }
 
 func (a *Alias) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (PathSeq, error) {
-	subSeq, err := a.subIt.IterResourcesImpl(ctx, subject)
+	subSeq, err := ctx.IterResources(a.subIt, subject)
 	if err != nil {
 		return nil, err
 	}
@@ -133,6 +133,7 @@ func (a *Alias) Clone() Iterator {
 
 func (a *Alias) Explain() Explain {
 	return Explain{
+		Name:       "Alias",
 		Info:       "Alias(" + a.relation + ")",
 		SubExplain: []Explain{a.subIt.Explain()},
 	}

pkg/query/arrow.go

@@ -1,6 +1,8 @@
 package query
 
 import (
+	"github.com/authzed/spicedb/internal/caveats"
+	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
@@ -35,45 +37,82 @@ func (a *Arrow) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRel
 	// don't restructure the tree, but can affect the best way to evaluate the tree, sometimes dynamically.
 
 	return func(yield func(Path, error) bool) {
-		for _, resource := range resources {
-			subit, err := a.left.IterSubjectsImpl(ctx, resource)
+		ctx.TraceStep(a, "processing %d resources", len(resources))
+
+		totalResultPaths := 0
+		for resourceIdx, resource := range resources {
+			ctx.TraceStep(a, "processing resource %d: %s:%s", resourceIdx, resource.ObjectType, resource.ObjectID)
+
+			subit, err := ctx.IterSubjects(a.left, resource)
 			if err != nil {
 				yield(Path{}, err)
 				return
 			}
+
+			leftPathCount := 0
 			for path, err := range subit {
 				if err != nil {
 					yield(Path{}, err)
 					return
 				}
+				leftPathCount++
+
 				checkResources := []Object{GetObject(path.Subject)}
-				checkit, err := a.right.CheckImpl(ctx, checkResources, subject)
+				ctx.TraceStep(a, "checking right side for subject %s:%s", path.Subject.ObjectType, path.Subject.ObjectID)
+
+				checkit, err := ctx.Check(a.right, checkResources, subject)
 				if err != nil {
 					yield(Path{}, err)
 					return
 				}
+
+				rightPathCount := 0
 				for checkPath, err := range checkit {
 					if err != nil {
 						yield(Path{}, err)
 						return
 					}
+					rightPathCount++
+
+					// Combine caveats from both sides using Path-based approach
+					// For arrow operations (left->right), both conditions must be satisfied (AND logic)
+					var combinedCaveat *core.CaveatExpression
+					if path.Caveat != nil && checkPath.Caveat != nil {
+						// Both sides have caveats - create combined caveat expression
+						combinedCaveat = caveats.And(path.Caveat, checkPath.Caveat)
+					} else if path.Caveat != nil {
+						// Only left side has caveat
+						combinedCaveat = path.Caveat
+					} else if checkPath.Caveat != nil {
+						// Only right side has caveat
+						combinedCaveat = checkPath.Caveat
+					}
+					// else both are nil, combinedCaveat remains nil
 
 					// Create combined path with resource from left and subject from right
 					combinedPath := Path{
 						Resource:   path.Resource,
 						Relation:   path.Relation,
 						Subject:    checkPath.Subject,
-						Caveat:     checkPath.Caveat,
+						Caveat:     combinedCaveat,
 						Expiration: checkPath.Expiration,
 						Integrity:  checkPath.Integrity,
 						Metadata:   make(map[string]any),
 					}
+
+					totalResultPaths++
 					if !yield(combinedPath, nil) {
 						return
 					}
 				}
+
+				ctx.TraceStep(a, "right side returned %d paths for subject %s:%s", rightPathCount, path.Subject.ObjectType, path.Subject.ObjectID)
 			}
+
+			ctx.TraceStep(a, "left side returned %d paths for resource %s:%s", leftPathCount, resource.ObjectType, resource.ObjectID)
 		}
+
+		ctx.TraceStep(a, "arrow completed with %d total result paths", totalResultPaths)
 	}, nil
 }
 
@@ -94,6 +133,7 @@ func (a *Arrow) Clone() Iterator {
 
 func (a *Arrow) Explain() Explain {
 	return Explain{
+		Name:       "Arrow",
 		Info:       "Arrow",
 		SubExplain: []Explain{a.left.Explain(), a.right.Explain()},
 	}