refactor: create new queryopt package to target query plan optimizations (#2970)

Author: barakmich

CHANGELOG.md

@@ -4,6 +4,9 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
+### Added
+- New `pkg/query/queryopt` package for building optimizations into the query plan
+
 ### Changed
 - Updated CI so that Postgres tests run against v18 which is GA and not against v13 which is EOL (https://github.com/authzed/spicedb/pull/2926)
 - Added tracing to request validation (https://github.com/authzed/spicedb/pull/2950)

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/authzed/spicedb/pkg/datalayer"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/query"
+	"github.com/authzed/spicedb/pkg/query/queryopt"
 	"github.com/authzed/spicedb/pkg/schema/v2"
 	"github.com/authzed/spicedb/pkg/tuple"
 	"github.com/authzed/spicedb/pkg/validationfile"
@@ -129,15 +130,18 @@ func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
 									require := require.New(t)
 
 									rel := assertion.Relationship
-									it, err := query.BuildIteratorFromSchema(handle.schema, rel.Resource.ObjectType, rel.Resource.Relation)
+									co, err := query.BuildOutlineFromSchema(handle.schema, rel.Resource.ObjectType, rel.Resource.Relation)
 									require.NoError(err)
 
 									// Apply static optimizations if requested
 									if optimizationMode.optimize {
-										it, _, err = query.ApplyOptimizations(it, query.StaticOptimizations)
+										co, err = queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations)
 										require.NoError(err)
 									}
 
+									it, err := co.Compile()
+									require.NoError(err)
+
 									qctx := handle.buildContext(t)
 
 									// Add caveat context from assertion if available
@@ -200,7 +204,9 @@ func runQueryPlanLookupResources(t *testing.T, handle *queryPlanConsistencyHandl
 				t.Run(tuple.StringONR(subject), func(t *testing.T) {
 					accessibleResources := accessibilitySet.LookupAccessibleResources(resourceRelation, subject)
 					queryCtx := handle.buildContext(t)
-					it, err := query.BuildIteratorFromSchema(handle.schema, resourceRelation.ObjectType, resourceRelation.Relation)
+					co, err := query.BuildOutlineFromSchema(handle.schema, resourceRelation.ObjectType, resourceRelation.Relation)
+					require.NoError(t, err)
+					it, err := co.Compile()
 					require.NoError(t, err)
 
 					// Perform a lookup call and ensure it returns the at least the same set of object IDs.
@@ -254,7 +260,9 @@ func runQueryPlanLookupSubjects(t *testing.T, handle *queryPlanConsistencyHandle
 				t.Run(tuple.StringONR(resource), func(t *testing.T) {
 					accessibleSubjects := accessibilitySet.LookupAccessibleSubjects(resource)
 					queryCtx := handle.buildContext(t)
-					it, err := query.BuildIteratorFromSchema(handle.schema, resourceRelation.ObjectType, resourceRelation.Relation)
+					co, err := query.BuildOutlineFromSchema(handle.schema, resourceRelation.ObjectType, resourceRelation.Relation)
+					require.NoError(t, err)
+					it, err := co.Compile()
 					require.NoError(t, err)
 
 					// Perform a lookup call and ensure it returns the at least the same set of subject IDs.

internal/services/v1/permissions_queryplan.go

@@ -10,6 +10,7 @@ import (
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/middleware/consistency"
 	"github.com/authzed/spicedb/pkg/query"
+	"github.com/authzed/spicedb/pkg/query/queryopt"
 	"github.com/authzed/spicedb/pkg/schema/v2"
 )
 
@@ -50,15 +51,19 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
 		return nil, ps.rewriteError(ctx, err)
 	}
 
-	// Build iterator tree from schema
-	// TODO: Better iterator caching
-	it, err := query.BuildIteratorFromSchema(fullSchema, req.Resource.ObjectType, req.Permission)
+	// Build and optimize the outline, then compile to an iterator tree.
+	// TODO: Better outline caching
+	co, err := query.BuildOutlineFromSchema(fullSchema, req.Resource.ObjectType, req.Permission)
 	if err != nil {
 		return nil, ps.rewriteError(ctx, err)
 	}
 
-	// Apply basic optimizations to the iterator tree
-	it, _, err = query.ApplyOptimizations(it, query.StaticOptimizations)
+	optimized, err := queryopt.ApplyOptimizations(co, queryopt.StandardOptimzations)
+	if err != nil {
+		return nil, ps.rewriteError(ctx, err)
+	}
+
+	it, err := optimized.Compile()
 	if err != nil {
 		return nil, ps.rewriteError(ctx, err)
 	}

pkg/query/advisor.go

@@ -62,8 +62,11 @@ func ApplyAdvisor(co CanonicalOutline, advisor PlanAdvisor) (CanonicalOutline, e
 // For each node it calls GetMutations on the advisor and applies the returned
 // mutations in sequence. After each mutation it verifies that the resulting
 // node's ID matches the original node's ID; a mismatch is a programmer bug.
+// After all mutations are applied, any newly synthesized nodes (ID==0) receive
+// fresh IDs via FillMissingNodeIDs, and their CanonicalKeys are recorded in
+// the provided keys map.
 func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdvisor) (Outline, error) {
-	return WalkOutlineBottomUp(outline, func(node Outline) (Outline, error) {
+	mutated, err := WalkOutlineBottomUp(outline, func(node Outline) (Outline, error) {
 		mutations, err := advisor.GetMutations(node, co)
 		if err != nil {
 			return Outline{}, err
@@ -81,6 +84,10 @@ func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdv
 		}
 		return result, nil
 	})
+	if err != nil {
+		return Outline{}, err
+	}
+	return FillMissingNodeIDs(mutated, co.CanonicalKeys), nil
 }
 
 // collectAdvisorHints walks the outline tree pre-order via WalkOutlinePreOrder,

pkg/query/build_tree.go

@@ -22,16 +22,6 @@ type outlineBuilder struct {
 	recursiveSentinels []*recursiveSentinelInfo     // Track recursion points for wrapping in RecursiveIterator
 }
 
-// BuildIteratorFromSchema takes a schema and walks the schema tree for a given definition namespace and a relationship or
-// permission therein. From this, it generates an iterator tree, rooted on that relationship.
-func BuildIteratorFromSchema(fullSchema *schema.Schema, definitionName string, relationName string) (Iterator, error) {
-	canonical, err := BuildOutlineFromSchema(fullSchema, definitionName, relationName)
-	if err != nil {
-		return nil, err
-	}
-	return canonical.Compile()
-}
-
 // BuildOutlineFromSchema builds a canonical Outline tree from the schema.
 func BuildOutlineFromSchema(fullSchema *schema.Schema, definitionName string, relationName string) (CanonicalOutline, error) {
 	builder := &outlineBuilder{

pkg/query/build_tree_test.go

@@ -14,6 +14,19 @@ import (
 	"github.com/authzed/spicedb/pkg/schema/v2"
 )
 
+// buildIterator is a test helper that builds a CanonicalOutline from the schema
+// and compiles it into an Iterator. It mirrors the old BuildIteratorFromSchema
+// convenience function which has been removed in favour of the explicit
+// BuildOutlineFromSchema → Compile() pipeline.
+func buildIterator(t *testing.T, fullSchema *schema.Schema, defName, relName string) (Iterator, error) {
+	t.Helper()
+	co, err := BuildOutlineFromSchema(fullSchema, defName, relName)
+	if err != nil {
+		return nil, err
+	}
+	return co.Compile()
+}
+
 func TestBuildTree(t *testing.T) {
 	t.Parallel()
 
@@ -28,7 +41,7 @@ func TestBuildTree(t *testing.T) {
 	dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 	require.NoError(err)
 
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "edit")
+	it, err := buildIterator(t, dsSchema, "document", "edit")
 	require.NoError(err)
 
 	ctx := NewLocalContext(t.Context(),
@@ -55,7 +68,7 @@ func TestBuildTreeMultipleRelations(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for edit permission which creates a union
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "edit")
+	it, err := buildIterator(t, dsSchema, "document", "edit")
 	require.NoError(err)
 
 	explain := it.Explain()
@@ -81,12 +94,12 @@ func TestBuildTreeInvalidDefinition(t *testing.T) {
 	require.NoError(err)
 
 	// Test with invalid definition name
-	_, err = BuildIteratorFromSchema(dsSchema, "nonexistent", "edit")
+	_, err = buildIterator(t, dsSchema, "nonexistent", "edit")
 	require.Error(err)
 	require.Contains(err.Error(), "couldn't find a schema definition named `nonexistent`")
 
 	// Test with invalid relation/permission name
-	_, err = BuildIteratorFromSchema(dsSchema, "document", "nonexistent")
+	_, err = buildIterator(t, dsSchema, "document", "nonexistent")
 	require.Error(err)
 	require.Contains(err.Error(), "couldn't find a relation or permission named `nonexistent`")
 }
@@ -105,7 +118,7 @@ func TestBuildTreeSubRelations(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for a relation with subrelations
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "parent")
+	it, err := buildIterator(t, dsSchema, "document", "parent")
 	require.NoError(err)
 
 	// Should have created a relation iterator
@@ -152,7 +165,7 @@ func TestBuildTreeRecursion(t *testing.T) {
 
 	// This should detect recursion and create a RecursiveIterator
 	// The arrow operation parent->member creates recursion: group->parent->member->parent->member...
-	it, err := BuildIteratorFromSchema(dsSchema, "group", "member")
+	it, err := buildIterator(t, dsSchema, "group", "member")
 	require.NoError(err)
 	require.NotNil(it)
 
@@ -203,7 +216,7 @@ func TestBuildTreeIntersectionOperation(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for view_and_edit permission which uses intersection operations
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "view_and_edit")
+	it, err := buildIterator(t, dsSchema, "document", "view_and_edit")
 	require.NoError(err)
 
 	// Should create an intersection iterator
@@ -245,7 +258,7 @@ func TestBuildTreeExclusionOperation(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for exclusion permission - should succeed
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "excluded_perm")
+	it, err := buildIterator(t, dsSchema, "document", "excluded_perm")
 	require.NoError(err)
 	require.NotNil(it)
 	// Should be wrapped in an Alias
@@ -297,7 +310,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "can_view")
+		it, err := buildIterator(t, dsSchema, "document", "can_view")
 		require.NoError(err)
 		require.NotNil(it)
 		// Should be wrapped in an Alias
@@ -336,7 +349,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "restricted_viewers")
+		it, err := buildIterator(t, dsSchema, "document", "restricted_viewers")
 		require.NoError(err)
 		require.NotNil(it)
 		// Should be wrapped in an Alias
@@ -377,7 +390,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "restricted_view")
+		it, err := buildIterator(t, dsSchema, "document", "restricted_view")
 		require.NoError(err)
 		require.NotNil(it)
 		// Should be wrapped in an Alias
@@ -419,7 +432,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "allowed_users")
+		it, err := buildIterator(t, dsSchema, "document", "allowed_users")
 		require.NoError(err)
 		require.NotNil(it)
 		// Should be wrapped in an Alias
@@ -457,7 +470,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		require.NoError(err)
 
 		// Building iterator should fail due to missing relation
-		_, err = BuildIteratorFromSchema(dsSchema, "document", "bad_exclusion")
+		_, err = buildIterator(t, dsSchema, "document", "bad_exclusion")
 		require.Error(err)
 		require.Contains(err.Error(), "couldn't find a relation or permission named `nonexistent_relation`")
 	})
@@ -480,7 +493,7 @@ func TestBuildTreeExclusionEdgeCases(t *testing.T) {
 		require.NoError(err)
 
 		// Building iterator should fail due to missing relation
-		_, err = BuildIteratorFromSchema(dsSchema, "document", "bad_exclusion")
+		_, err = buildIterator(t, dsSchema, "document", "bad_exclusion")
 		require.Error(err)
 		require.Contains(err.Error(), "couldn't find a relation or permission named `nonexistent_relation`")
 	})
@@ -507,7 +520,7 @@ func TestBuildTreeArrowMissingLeftRelation(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for arrow with missing left relation
-	_, err = BuildIteratorFromSchema(dsSchema, "document", "bad_arrow")
+	_, err = buildIterator(t, dsSchema, "document", "bad_arrow")
 	require.Error(err)
 	require.Contains(err.Error(), "couldn't find left-hand relation for arrow")
 }
@@ -526,7 +539,7 @@ func TestBuildTreeSingleRelationOptimization(t *testing.T) {
 	require.NoError(err)
 
 	// Test building iterator for a simple relation - should not create unnecessary unions
-	it, err := BuildIteratorFromSchema(dsSchema, "document", "owner")
+	it, err := buildIterator(t, dsSchema, "document", "owner")
 	require.NoError(err)
 
 	// Should create a simple relation iterator without extra union wrappers
@@ -580,7 +593,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.NoError(err)
 
 		// Should create an alias wrapping union with arrow
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -617,7 +630,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -649,7 +662,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.NoError(err)
 
 		// Should create RecursiveIterator for arrow recursion
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -680,7 +693,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		require.NoError(err)
 
 		// Should fail when trying to build iterator due to missing subrelation
-		_, err = BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		_, err = buildIterator(t, dsSchema, "document", "viewer")
 		require.Error(err)
 		require.Contains(err.Error(), "couldn't find a relation or permission named `nonexistent`")
 	})
@@ -708,7 +721,7 @@ func TestBuildTreeSubrelationHandling(t *testing.T) {
 		dsSchema, err := schema.BuildSchemaFromDefinitions(objectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -770,7 +783,7 @@ func TestBuildTreeWildcardIterator(t *testing.T) {
 
 	t.Run("Schema with wildcard creates WildcardIterator", func(t *testing.T) {
 		t.Parallel()
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -801,7 +814,7 @@ func TestBuildTreeWildcardIterator(t *testing.T) {
 		mixedSchema, err := schema.BuildSchemaFromDefinitions(mixedObjectDefs, nil)
 		require.NoError(err)
 
-		it, err := BuildIteratorFromSchema(mixedSchema, "document", "viewer")
+		it, err := buildIterator(t, mixedSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -859,7 +872,7 @@ func TestBuildTreeMutualRecursionSentinelFiltering(t *testing.T) {
 	t.Run("document viewer builds successfully with mutual recursion", func(t *testing.T) {
 		t.Parallel()
 		// Build iterator for document#viewer - should detect recursion and wrap properly
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -872,7 +885,7 @@ func TestBuildTreeMutualRecursionSentinelFiltering(t *testing.T) {
 	t.Run("otherdocument viewer builds successfully with mutual recursion", func(t *testing.T) {
 		t.Parallel()
 		// Build iterator for otherdocument#viewer - should also handle mutual recursion
-		it, err := BuildIteratorFromSchema(dsSchema, "otherdocument", "viewer")
+		it, err := buildIterator(t, dsSchema, "otherdocument", "viewer")
 		require.NoError(err)
 		require.NotNil(it)
 
@@ -890,7 +903,7 @@ func TestBuildTreeMutualRecursionSentinelFiltering(t *testing.T) {
 		// its own sentinels.
 
 		// Build the tree
-		it, err := BuildIteratorFromSchema(dsSchema, "document", "viewer")
+		it, err := buildIterator(t, dsSchema, "document", "viewer")
 		require.NoError(err)
 		require.NotNil(it)