Merge pull request #2550 from authzed/barakmich/refactor_objects

refactor: Query iterators now operate in Objects and ObjectRelations instead of strings

Author: barakmich

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -119,7 +119,7 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 
 							qctx := handle.buildContext(t)
 
-							seq, err := qctx.Check(it, []string{rel.Resource.ObjectID}, rel.Subject.ObjectID)
+							seq, err := qctx.Check(it, []query.Object{query.GetObject(rel.Resource)}, rel.Subject)
 							require.NoError(err)
 
 							rels, err := query.CollectAll(seq)

pkg/query/arrow.go

@@ -23,7 +23,7 @@ func NewArrow(left, right Iterator) *Arrow {
 	}
 }
 
-func (a *Arrow) CheckImpl(ctx *Context, resourceIDs []string, subjectID string) (RelationSeq, error) {
+func (a *Arrow) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
 	// TODO -- the ordering, directionality, batching, everything can depend on other statistics.
 	//
 	// There are three major strategies:
@@ -36,8 +36,8 @@ func (a *Arrow) CheckImpl(ctx *Context, resourceIDs []string, subjectID string)
 	// don't restructure the tree, but can affect the best way to evaluate the tree, sometimes dynamically.
 
 	return func(yield func(Relation, error) bool) {
-		for _, rid := range resourceIDs {
-			subit, err := ctx.IterSubjects(a.left, rid)
+		for _, resource := range resources {
+			subit, err := a.left.IterSubjectsImpl(ctx, resource)
 			if err != nil {
 				yield(Relation{}, err)
 				return
@@ -47,7 +47,8 @@ func (a *Arrow) CheckImpl(ctx *Context, resourceIDs []string, subjectID string)
 					yield(Relation{}, err)
 					return
 				}
-				checkit, err := ctx.Check(a.right, []string{rel.Subject.ObjectID}, subjectID)
+				checkResources := []Object{GetObject(rel.Subject)}
+				checkit, err := a.right.CheckImpl(ctx, checkResources, subject)
 				if err != nil {
 					yield(Relation{}, err)
 					return
@@ -75,11 +76,11 @@ func (a *Arrow) CheckImpl(ctx *Context, resourceIDs []string, subjectID string)
 	}, nil
 }
 
-func (a *Arrow) IterSubjectsImpl(ctx *Context, resourceID string) (RelationSeq, error) {
+func (a *Arrow) IterSubjectsImpl(ctx *Context, resource Object) (RelationSeq, error) {
 	return nil, spiceerrors.MustBugf("unimplemented")
 }
 
-func (a *Arrow) IterResourcesImpl(ctx *Context, subjectID string) (RelationSeq, error) {
+func (a *Arrow) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (RelationSeq, error) {
 	return nil, spiceerrors.MustBugf("unimplemented")
 }
 

pkg/query/arrow_test.go

@@ -33,7 +33,7 @@ func TestArrowIterator(t *testing.T) {
 
 		// Test arrow operation: find resources where left side connects to right side
 		// This looks for documents whose parent folder has viewers
-		relSeq, err := ctx.Check(arrow, []string{"spec1", "spec2"}, "alice")
+		relSeq, err := ctx.Check(arrow, NewObjects("document", "spec1", "spec2"), NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 
 		rels, err := CollectAll(relSeq)
@@ -61,7 +61,7 @@ func TestArrowIterator(t *testing.T) {
 			Executor: LocalExecutor{},
 		}
 
-		relSeq, err := ctx.Check(arrow, []string{}, "alice")
+		relSeq, err := ctx.Check(arrow, []Object{}, NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 
 		rels, err := CollectAll(relSeq)
@@ -78,7 +78,7 @@ func TestArrowIterator(t *testing.T) {
 			Executor: LocalExecutor{},
 		}
 
-		relSeq, err := ctx.Check(arrow, []string{"nonexistent"}, "alice")
+		relSeq, err := ctx.Check(arrow, NewObjects("document", "nonexistent"), NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 
 		rels, err := CollectAll(relSeq)
@@ -96,7 +96,7 @@ func TestArrowIterator(t *testing.T) {
 			Executor: LocalExecutor{},
 		}
 
-		relSeq, err := ctx.Check(arrow, []string{"spec1"}, "nonexistent")
+		relSeq, err := ctx.Check(arrow, NewObjects("document", "spec1"), NewObject("user", "nonexistent").WithEllipses())
 		require.NoError(err)
 
 		rels, err := CollectAll(relSeq)
@@ -115,7 +115,7 @@ func TestArrowIterator(t *testing.T) {
 		}
 
 		require.Panics(func() {
-			_, _ = ctx.IterSubjects(arrow, "spec1")
+			_, _ = ctx.IterSubjects(arrow, NewObject("document", "spec1"))
 		})
 	})
 
@@ -129,7 +129,7 @@ func TestArrowIterator(t *testing.T) {
 		}
 
 		require.Panics(func() {
-			_, _ = ctx.IterResources(arrow, "alice")
+			_, _ = ctx.IterResources(arrow, NewObject("user", "alice").WithEllipses())
 		})
 	})
 }
@@ -164,13 +164,13 @@ func TestArrowIteratorClone(t *testing.T) {
 	subjectID := "alice"
 
 	// Collect results from original iterator
-	originalSeq, err := ctx.Check(original, resourceIDs, subjectID)
+	originalSeq, err := ctx.Check(original, NewObjects("document", resourceIDs[0]), NewObject("user", subjectID).WithEllipses())
 	require.NoError(err)
 	originalResults, err := CollectAll(originalSeq)
 	require.NoError(err)
 
 	// Collect results from cloned iterator
-	clonedSeq, err := ctx.Check(cloned, resourceIDs, subjectID)
+	clonedSeq, err := ctx.Check(cloned, NewObjects("document", resourceIDs[0]), NewObject("user", subjectID).WithEllipses())
 	require.NoError(err)
 	clonedResults, err := CollectAll(clonedSeq)
 	require.NoError(err)
@@ -213,7 +213,7 @@ func TestArrowIteratorMultipleResources(t *testing.T) {
 	}
 
 	// Test with multiple resource IDs
-	relSeq, err := ctx.Check(arrow, []string{"spec1", "spec2", "nonexistent"}, "alice")
+	relSeq, err := ctx.Check(arrow, NewObjects("document", "spec1", "spec2", "nonexistent"), NewObject("user", "alice").WithEllipses())
 	require.NoError(err)
 
 	rels, err := CollectAll(relSeq)

pkg/query/build_tree_test.go

@@ -37,7 +37,7 @@ func TestBuildTree(t *testing.T) {
 		Revision:  revision,
 	}
 
-	relSeq, err := ctx.Check(it, []string{"specialplan"}, "multiroleguy")
+	relSeq, err := ctx.Check(it, NewObjects("document", "specialplan"), NewObject("user", "multiroleguy").WithEllipses())
 	require.NoError(err)
 
 	_, err = CollectAll(relSeq)
@@ -71,7 +71,7 @@ func TestBuildTreeMultipleRelations(t *testing.T) {
 		Revision:  revision,
 	}
 
-	relSeq, err := ctx.Check(it, []string{"specialplan"}, "multiroleguy")
+	relSeq, err := ctx.Check(it, NewObjects("document", "specialplan"), NewObject("user", "multiroleguy").WithEllipses())
 	require.NoError(err)
 
 	rels, err := CollectAll(relSeq)
@@ -127,7 +127,7 @@ func TestBuildTreeSubRelations(t *testing.T) {
 	}
 
 	// Just test that the iterator can be executed without error
-	relSeq, err := ctx.Check(it, []string{"companyplan"}, "legal")
+	relSeq, err := ctx.Check(it, NewObjects("document", "companyplan"), NewObject("user", "legal").WithEllipses())
 	require.NoError(err)
 
 	_, err = CollectAll(relSeq)
@@ -214,7 +214,7 @@ func TestBuildTreeIntersectionOperation(t *testing.T) {
 	}
 
 	// Test execution
-	relSeq, err := ctx.Check(it, []string{"specialplan"}, "multiroleguy")
+	relSeq, err := ctx.Check(it, NewObjects("document", "specialplan"), NewObject("user", "multiroleguy").WithEllipses())
 	require.NoError(err)
 
 	_, err = CollectAll(relSeq)
@@ -301,7 +301,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		require.IsType(&Exclusion{}, it)
 
 		// Test execution doesn't crash
-		relSeq, err := ctx.Check(it, []string{"test_doc"}, "alice")
+		relSeq, err := ctx.Check(it, []Object{NewObject("document", "test_doc")}, NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 		_, err = CollectAll(relSeq)
 		require.NoError(err)
@@ -522,7 +522,7 @@ func TestBuildTreeSingleRelationOptimization(t *testing.T) {
 	}
 
 	// Test execution
-	relSeq, err := ctx.Check(it, []string{"companyplan"}, "legal")
+	relSeq, err := ctx.Check(it, NewObjects("document", "companyplan"), NewObject("user", "legal").WithEllipses())
 	require.NoError(err)
 
 	_, err = CollectAll(relSeq)
@@ -607,7 +607,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.Contains(explainStr, "Union") // Should contain union for base relation + arrow
 
 		// Test execution doesn't crash
-		relSeq, err := ctx.Check(it, []string{"test_doc"}, "alice")
+		relSeq, err := ctx.Check(it, []Object{NewObject("document", "test_doc")}, NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 		_, err = CollectAll(relSeq)
 		require.NoError(err)
@@ -695,7 +695,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.Contains(explainStr, "Union") // Should contain union for different relation types
 
 		// Test execution doesn't crash
-		relSeq, err := ctx.Check(it, []string{"test_doc"}, "alice")
+		relSeq, err := ctx.Check(it, []Object{NewObject("document", "test_doc")}, NewObject("user", "alice").WithEllipses())
 		require.NoError(err)
 		_, err = CollectAll(relSeq)
 		require.NoError(err)

pkg/query/context.go

@@ -19,43 +19,43 @@ type Context struct {
 }
 
 // Check tests if, for the underlying set of relationships (which may be a full expression or a basic lookup, depending on the iterator)
-// any of the `resourceIDs` are connected to `subjectID`.
-// Returns the sequence of matching relations, if they exist, at most `len(resourceIDs)`.
-func (ctx *Context) Check(it Iterator, resourceIDs []string, subjectID string) (RelationSeq, error) {
+// any of the `resources` are connected to `subject`.
+// Returns the sequence of matching relations, if they exist, at most `len(resources)`.
+func (ctx *Context) Check(it Iterator, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
 	if ctx.Executor == nil {
 		return nil, spiceerrors.MustBugf("no executor has been set")
 	}
-	return ctx.Executor.Check(ctx, it, resourceIDs, subjectID)
+	return ctx.Executor.Check(ctx, it, resources, subject)
 }
 
-// IterSubjects returns a sequence of all the relations in this set that match the given resourceID.
-func (ctx *Context) IterSubjects(it Iterator, resourceID string) (RelationSeq, error) {
+// IterSubjects returns a sequence of all the relations in this set that match the given resource.
+func (ctx *Context) IterSubjects(it Iterator, resource Object) (RelationSeq, error) {
 	if ctx.Executor == nil {
 		return nil, spiceerrors.MustBugf("no executor has been set")
 	}
-	return ctx.Executor.IterSubjects(ctx, it, resourceID)
+	return ctx.Executor.IterSubjects(ctx, it, resource)
 }
 
-// IterResources returns a sequence of all the relations in this set that match the given subjectID.
-func (ctx *Context) IterResources(it Iterator, subjectID string) (RelationSeq, error) {
+// IterResources returns a sequence of all the relations in this set that match the given subject.
+func (ctx *Context) IterResources(it Iterator, subject ObjectAndRelation) (RelationSeq, error) {
 	if ctx.Executor == nil {
 		return nil, spiceerrors.MustBugf("no executor has been set")
 	}
-	return ctx.Executor.IterResources(ctx, it, subjectID)
+	return ctx.Executor.IterResources(ctx, it, subject)
 }
 
 // Executor as chooses how to proceed given an iterator -- perhaps in parallel, perhaps by RPC, etc -- and chooses how to process iteration from the subtree.
 // The correctness logic for the results that are generated are up to each iterator, and each iterator may use statistics to choose the best, yet still correct, logical evaluation strategy.
 // The Executor, meanwhile, makes that evaluation happen in the most convienent form, based on its implementation.
 type Executor interface {
 	// Check tests if, for the underlying set of relationships (which may be a full expression or a basic lookup, depending on the iterator)
-	// any of the `resourceIDs` are connected to `subjectID`.
-	// Returns the sequence of matching relations, if they exist, at most `len(resourceIDs)`.
-	Check(ctx *Context, it Iterator, resourceIDs []string, subjectID string) (RelationSeq, error)
+	// any of the `resources` are connected to `subject`.
+	// Returns the sequence of matching relations, if they exist, at most `len(resources)`.
+	Check(ctx *Context, it Iterator, resources []Object, subject ObjectAndRelation) (RelationSeq, error)
 
-	// IterSubjects returns a sequence of all the relations in this set that match the given resourceID.
-	IterSubjects(ctx *Context, it Iterator, resourceID string) (RelationSeq, error)
+	// IterSubjects returns a sequence of all the relations in this set that match the given resource.
+	IterSubjects(ctx *Context, it Iterator, resource Object) (RelationSeq, error)
 
-	// IterResources returns a sequence of all the relations in this set that match the given subjectID.
-	IterResources(ctx *Context, it Iterator, subjectID string) (RelationSeq, error)
+	// IterResources returns a sequence of all the relations in this set that match the given subject.
+	IterResources(ctx *Context, it Iterator, subject ObjectAndRelation) (RelationSeq, error)
 }

pkg/query/datastore.go

@@ -35,15 +35,20 @@ func (r *RelationIterator) buildSubjectRelationFilter() datastore.SubjectRelatio
 	return datastore.SubjectRelationFilter{}.WithNonEllipsisRelation(r.base.Subrelation)
 }
 
-func (r *RelationIterator) CheckImpl(ctx *Context, resourceIDs []string, subjectID string) (RelationSeq, error) {
+func (r *RelationIterator) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
+	resourceIDs := make([]string, len(resources))
+	for i, res := range resources {
+		resourceIDs[i] = res.ObjectID
+	}
+
 	filter := datastore.RelationshipsFilter{
 		OptionalResourceType:     r.base.DefinitionName(),
 		OptionalResourceIds:      resourceIDs,
 		OptionalResourceRelation: r.base.RelationName(),
 		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
 			{
 				OptionalSubjectType: r.base.Type,
-				OptionalSubjectIds:  []string{subjectID},
+				OptionalSubjectIds:  []string{subject.ObjectID},
 				RelationFilter:      r.buildSubjectRelationFilter(),
 			},
 		},
@@ -63,10 +68,10 @@ func (r *RelationIterator) CheckImpl(ctx *Context, resourceIDs []string, subject
 	return RelationSeq(relIter), nil
 }
 
-func (r *RelationIterator) IterSubjectsImpl(ctx *Context, resourceID string) (RelationSeq, error) {
+func (r *RelationIterator) IterSubjectsImpl(ctx *Context, resource Object) (RelationSeq, error) {
 	filter := datastore.RelationshipsFilter{
 		OptionalResourceType:     r.base.DefinitionName(),
-		OptionalResourceIds:      []string{resourceID},
+		OptionalResourceIds:      []string{resource.ObjectID},
 		OptionalResourceRelation: r.base.RelationName(),
 		OptionalSubjectsSelectors: []datastore.SubjectsSelector{
 			{
@@ -90,7 +95,7 @@ func (r *RelationIterator) IterSubjectsImpl(ctx *Context, resourceID string) (Re
 	return RelationSeq(relIter), nil
 }
 
-func (r *RelationIterator) IterResourcesImpl(ctx *Context, subjectID string) (RelationSeq, error) {
+func (r *RelationIterator) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (RelationSeq, error) {
 	return nil, spiceerrors.MustBugf("unimplemented")
 }
 

pkg/query/exclusion.go

@@ -20,9 +20,9 @@ func NewExclusion(mainSet, excluded Iterator) *Exclusion {
 	}
 }
 
-func (e *Exclusion) CheckImpl(ctx *Context, resourceIDs []string, subjectID string) (RelationSeq, error) {
+func (e *Exclusion) CheckImpl(ctx *Context, resources []Object, subject ObjectAndRelation) (RelationSeq, error) {
 	// Get all relations from the main set
-	mainSeq, err := ctx.Check(e.mainSet, resourceIDs, subjectID)
+	mainSeq, err := e.mainSet.CheckImpl(ctx, resources, subject)
 	if err != nil {
 		return nil, err
 	}
@@ -40,7 +40,7 @@ func (e *Exclusion) CheckImpl(ctx *Context, resourceIDs []string, subjectID stri
 	}
 
 	// Get all relations from the excluded set
-	excludedSeq, err := ctx.Check(e.excluded, resourceIDs, subjectID)
+	excludedSeq, err := e.excluded.CheckImpl(ctx, resources, subject)
 	if err != nil {
 		return nil, err
 	}
@@ -57,10 +57,8 @@ func (e *Exclusion) CheckImpl(ctx *Context, resourceIDs []string, subjectID stri
 
 			// Check if this relation exists in the excluded set (only compare endpoints: resource and subject object types/IDs)
 			for _, excludedRel := range excludedRels {
-				if mainRel.Resource.ObjectType == excludedRel.Resource.ObjectType &&
-					mainRel.Resource.ObjectID == excludedRel.Resource.ObjectID &&
-					mainRel.Subject.ObjectType == excludedRel.Subject.ObjectType &&
-					mainRel.Subject.ObjectID == excludedRel.Subject.ObjectID {
+				if GetObject(mainRel.Resource).Equals(GetObject(excludedRel.Resource)) &&
+					GetObject(mainRel.Subject).Equals(GetObject(excludedRel.Subject)) {
 					found = true
 					break
 				}
@@ -76,11 +74,11 @@ func (e *Exclusion) CheckImpl(ctx *Context, resourceIDs []string, subjectID stri
 	}, nil
 }
 
-func (e *Exclusion) IterSubjectsImpl(ctx *Context, resourceID string) (RelationSeq, error) {
+func (e *Exclusion) IterSubjectsImpl(ctx *Context, resource Object) (RelationSeq, error) {
 	return nil, spiceerrors.MustBugf("unimplemented")
 }
 
-func (e *Exclusion) IterResourcesImpl(ctx *Context, subjectID string) (RelationSeq, error) {
+func (e *Exclusion) IterResourcesImpl(ctx *Context, subject ObjectAndRelation) (RelationSeq, error) {
 	return nil, spiceerrors.MustBugf("unimplemented")
 }