refactor(query): add QueryDatastoreReader to act as an implementation shim for datastores

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -48,7 +48,7 @@ type queryPlanConsistencyHandle struct {
 
 func (q *queryPlanConsistencyHandle) buildContext(t *testing.T) *query.Context {
 	return query.NewLocalContext(t.Context(),
-		query.WithReader(datalayer.NewDataLayer(q.ds).SnapshotReader(q.revision)),
+		query.WithRevisionedReader(datalayer.NewDataLayer(q.ds).SnapshotReader(q.revision)),
 		query.WithCaveatRunner(caveats.NewCaveatRunner(caveattypes.Default.TypeSet)),
 		query.WithTraceLogger(query.NewTraceLogger())) // Enable tracing for debugging
 }

internal/services/v1/permissions_queryplan.go

@@ -73,7 +73,7 @@
 	qctx := &query.Context{
 		Context:       ctx,
 		Executor:      query.LocalExecutor{},
-		Reader:        reader,
+		Reader:        query.NewQueryDatastoreReader(reader),
 		CaveatContext: caveatContext,
 		CaveatRunner:  caveatsimpl.NewCaveatRunner(ps.config.CaveatTypeSet),
 	}

pkg/query/alias.go

@@ -1,10 +1,5 @@
 package query
 
-import (
-	"github.com/authzed/spicedb/pkg/datastore"
-	"github.com/authzed/spicedb/pkg/datastore/options"
-)
-
 // AliasIterator is an iterator that rewrites the Resource's Relation field of all paths
 // streamed from the sub-iterator to a specified alias relation.
 type AliasIterator struct {
@@ -144,31 +139,7 @@ func (a *AliasIterator) shouldIncludeSelfEdge(ctx *Context, resource Object, fil
 // resourceExistsAsSubject queries the datastore to check if the given resource appears
 // as a subject in any relationship, including expired relationships.
 func (a *AliasIterator) resourceExistsAsSubject(ctx *Context, resource Object) (bool, error) {
-	filter := datastore.RelationshipsFilter{
-		OptionalSubjectsSelectors: []datastore.SubjectsSelector{{
-			OptionalSubjectType: resource.ObjectType,
-			OptionalSubjectIds:  []string{resource.ObjectID},
-			RelationFilter:      datastore.SubjectRelationFilter{}.WithNonEllipsisRelation(a.relation),
-		}},
-		OptionalExpirationOption: datastore.ExpirationFilterOptionNone,
-	}
-
-	iter, err := ctx.Reader.QueryRelationships(ctx, filter,
-		options.WithLimit(options.LimitOne),
-		options.WithSkipExpiration(true)) // Include expired relationships
-	if err != nil {
-		return false, err
-	}
-
-	// Check if any relationship exists
-	for _, err := range iter {
-		if err != nil {
-			return false, err
-		}
-		return true, nil
-	}
-
-	return false, nil
+	return ctx.Reader.SubjectExistsAsRelationship(ctx, resource, a.relation)
 }
 
 func (a *AliasIterator) IterResourcesImpl(ctx *Context, subject ObjectAndRelation, filterResourceType ObjectType) (PathSeq, error) {

pkg/query/arrow_reversal_test.go

@@ -112,15 +112,15 @@ func TestDoubleWideArrowAdvisedMatchesPlain(t *testing.T) {
 	resources := NewObjects("file", "file0")
 	subject := NewObject("user", "user42").WithEllipses()
 
-	reader := datalayer.NewDataLayer(rawDS).SnapshotReader(revision)
+	readerOpt := WithRevisionedReader(datalayer.NewDataLayer(rawDS).SnapshotReader(revision))
 
 	// ---- plain (LTR) ----
 
 	plainTrace := NewTraceLogger()
 	plainIt, err := canonicalOutline.Compile()
 	require.NoError(t, err)
 
-	plainSeq, err := NewLocalContext(ctx, WithReader(reader), WithTraceLogger(plainTrace)).
+	plainSeq, err := NewLocalContext(ctx, readerOpt, WithTraceLogger(plainTrace)).
 		Check(plainIt, resources, subject)
 	require.NoError(t, err)
 	plainPaths, err := CollectAll(plainSeq)
@@ -132,7 +132,7 @@ func TestDoubleWideArrowAdvisedMatchesPlain(t *testing.T) {
 	obs := NewCountObserver()
 	warmIt, err := canonicalOutline.Compile()
 	require.NoError(t, err)
-	warmSeq, err := NewLocalContext(ctx, WithReader(reader), WithObserver(obs)).
+	warmSeq, err := NewLocalContext(ctx, readerOpt, WithObserver(obs)).
 		Check(warmIt, resources, subject)
 	require.NoError(t, err)
 	_, err = CollectAll(warmSeq)
@@ -144,7 +144,7 @@ func TestDoubleWideArrowAdvisedMatchesPlain(t *testing.T) {
 	require.NoError(t, err)
 
 	advisedTrace := NewTraceLogger()
-	advisedSeq, err := NewLocalContext(ctx, WithReader(reader), WithTraceLogger(advisedTrace)).
+	advisedSeq, err := NewLocalContext(ctx, readerOpt, WithTraceLogger(advisedTrace)).
 		Check(advisedIt, resources, subject)
 	require.NoError(t, err)
 	advisedPaths, err := CollectAll(advisedSeq)

pkg/query/benchmarks/check_deep_arrow_benchmark_test.go

@@ -20,14 +20,22 @@ import (
 
 // BenchmarkCheckDeepArrow benchmarks permission checking through a deep recursive chain.
 // This recreates the testharness scenario with:
-// - A 30+ level deep parent chain: document:target -> document:1 -> ... -> document:29
+// - A 30+ level deep parent chain: document:target -> document:1 -> ... -> document:30
 // - document:29#view@user:slow
 // - Checking if user:slow has viewer permission on document:target
 //
 // The permission viewer = view + parent->viewer creates a recursive traversal through
 // all 30+ levels to find the view relationship at the end of the chain.
+//
+// Four sub-benchmarks are run:
+//   - plain:         compile the outline directly and run Check each iteration
+//   - advised:       seed a CountAdvisor from a single warm-up run, apply it to the
+//     canonical outline, compile once, then run Check each iteration
+//   - plain_delay:   same as plain, but with a delay reader simulating network latency
+//   - advised_delay: same as advised, but with a delay reader simulating network latency
 func BenchmarkCheckDeepArrow(b *testing.B) {
-	// Create an in-memory datastore
+	// ---- shared setup ----
+
 	rawDS, err := memdb.NewMemdbDatastore(0, 0, memdb.DisableGC)
 	require.NoError(b, err)
 
@@ -43,76 +51,159 @@ func BenchmarkCheckDeepArrow(b *testing.B) {
 		}
 	`
 
-	// Compile the schema
 	compiled, err := compiler.Compile(compiler.InputSchema{
 		Source:       input.Source("benchmark"),
 		SchemaString: schemaText,
 	}, compiler.AllowUnprefixedObjectType())
 	require.NoError(b, err)
 
-	// Write the schema
 	_, err = rawDS.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
 		return rwt.LegacyWriteNamespaces(ctx, compiled.ObjectDefinitions...)
 	})
 	require.NoError(b, err)
 
-	// Build relationships for the deep arrow scenario
-	// Create a chain: document:target -> document:1 -> document:2 -> ... -> document:30 -> document:31
+	// Build relationships for the deep arrow scenario.
+	// Chain: document:target -> document:1 -> document:2 -> ... -> document:30
 	// Plus: document:29#view@user:slow
 	relationships := make([]tuple.Relationship, 0, 33)
-
-	// document:target#parent@document:1
 	relationships = append(relationships, tuple.MustParse("document:target#parent@document:1"))
-
-	// Chain: document:1 through document:30
 	for i := 1; i <= 30; i++ {
 		rel := fmt.Sprintf("document:%d#parent@document:%d", i, i+1)
 		relationships = append(relationships, tuple.MustParse(rel))
 	}
-
-	// The view relationship at the end of the chain
 	relationships = append(relationships, tuple.MustParse("document:29#view@user:slow"))
 
-	// Write all relationships to the datastore
 	revision, err := common.WriteRelationships(ctx, rawDS, tuple.UpdateOperationCreate, relationships...)
 	require.NoError(b, err)
 
-	// Build schema for querying
 	dsSchema, err := schema.BuildSchemaFromDefinitions(compiled.ObjectDefinitions, nil)
 	require.NoError(b, err)
 
-	// Create the iterator tree for the viewer permission using BuildIteratorFromSchema
-	viewerIterator, err := query.BuildIteratorFromSchema(dsSchema, "document", "viewer")
+	// Build the canonical outline once; all sub-benchmarks derive from it.
+	canonicalOutline, err := query.BuildOutlineFromSchema(dsSchema, "document", "viewer")
 	require.NoError(b, err)
 
-	// Create query context
-	queryCtx := query.NewLocalContext(ctx,
-		query.WithReader(datalayer.NewDataLayer(rawDS).SnapshotReader(revision)),
-		query.WithMaxRecursionDepth(50),
-	)
-
-	// The resource we're checking: document:target
+	// The resource and subject are the same for all sub-benchmarks.
 	resources := query.NewObjects("document", "target")
-
-	// The subject we're checking: user:slow
 	subject := query.NewObject("user", "slow").WithEllipses()
 
-	// Reset the timer - everything before this is setup
-	b.ResetTimer()
+	// Base reader (no simulated latency).
+	reader := query.NewQueryDatastoreReader(datalayer.NewDataLayer(rawDS).SnapshotReader(revision))
 
-	// Run the benchmark
-	for b.Loop() {
-		// Check if user:slow can view document:target
-		// This will traverse the entire 30+ level chain
-		seq, err := queryCtx.Check(viewerIterator, resources, subject)
-		require.NoError(b, err)
+	// Delay reader wrapping the base reader with simulated network latency.
+	delayReader := query.NewDelayReader(networkDelay, reader)
 
-		// Collect all results (should find user:slow at the end of the chain)
-		paths, err := query.CollectAll(seq)
+	// buildAdvisedIterator seeds a CountAdvisor from a single warm-up run using the
+	// provided reader and returns the compiled advised iterator.
+	buildAdvisedIterator := func(b *testing.B, r query.QueryDatastoreReader) query.Iterator {
+		b.Helper()
+		obs := query.NewCountObserver()
+		warmIt, err := canonicalOutline.Compile()
+		require.NoError(b, err)
+		warmCtx := query.NewLocalContext(ctx,
+			query.WithReader(r),
+			query.WithObserver(obs),
+			query.WithMaxRecursionDepth(50),
+		)
+		seq, err := warmCtx.Check(warmIt, resources, subject)
+		require.NoError(b, err)
+		_, err = query.CollectAll(seq)
 		require.NoError(b, err)
 
-		// Verify we found the expected result
-		require.Len(b, paths, 1)
-		require.Equal(b, "slow", paths[0].Subject.ObjectID)
+		advisor := query.NewCountAdvisor(obs.GetStats())
+		advisedCO, err := query.ApplyAdvisor(canonicalOutline, advisor)
+		require.NoError(b, err)
+		advisedIt, err := advisedCO.Compile()
+		require.NoError(b, err)
+		return advisedIt
 	}
+
+	// ---- plain sub-benchmark ----
+
+	b.Run("plain", func(b *testing.B) {
+		it, err := canonicalOutline.Compile()
+		require.NoError(b, err)
+
+		b.Log("plain explain:\n", it.Explain())
+
+		queryCtx := query.NewLocalContext(ctx,
+			query.WithReader(reader),
+			query.WithMaxRecursionDepth(50),
+		)
+
+		b.ResetTimer()
+		for b.Loop() {
+			seq, err := queryCtx.Check(it, resources, subject)
+			require.NoError(b, err)
+			paths, err := query.CollectAll(seq)
+			require.NoError(b, err)
+			require.Len(b, paths, 1)
+			require.Equal(b, "slow", paths[0].Subject.ObjectID)
+		}
+	})
+
+	// ---- advised sub-benchmark ----
+
+	b.Run("advised", func(b *testing.B) {
+		advisedIt := buildAdvisedIterator(b, reader)
+
+		b.Log("advised explain:\n", advisedIt.Explain())
+
+		queryCtx := query.NewLocalContext(ctx,
+			query.WithReader(reader),
+			query.WithMaxRecursionDepth(50),
+		)
+
+		b.ResetTimer()
+		for b.Loop() {
+			seq, err := queryCtx.Check(advisedIt, resources, subject)
+			require.NoError(b, err)
+			paths, err := query.CollectAll(seq)
+			require.NoError(b, err)
+			require.Len(b, paths, 1)
+			require.Equal(b, "slow", paths[0].Subject.ObjectID)
+		}
+	})
+
+	// ---- plain_delay sub-benchmark ----
+
+	b.Run("plain_delay", func(b *testing.B) {
+		it, err := canonicalOutline.Compile()
+		require.NoError(b, err)
+
+		queryCtx := query.NewLocalContext(ctx,
+			query.WithReader(delayReader),
+			query.WithMaxRecursionDepth(50),
+		)
+
+		b.ResetTimer()
+		for b.Loop() {
+			seq, err := queryCtx.Check(it, resources, subject)
+			require.NoError(b, err)
+			paths, err := query.CollectAll(seq)
+			require.NoError(b, err)
+			require.Len(b, paths, 1)
+			require.Equal(b, "slow", paths[0].Subject.ObjectID)
+		}
+	})
+
+	// ---- advised_delay sub-benchmark ----
+
+	b.Run("advised_delay", func(b *testing.B) {
+		advisedIt := buildAdvisedIterator(b, delayReader)
+		queryCtx := query.NewLocalContext(ctx,
+			query.WithReader(delayReader),
+			query.WithMaxRecursionDepth(50),
+		)
+
+		b.ResetTimer()
+		for b.Loop() {
+			seq, err := queryCtx.Check(advisedIt, resources, subject)
+			require.NoError(b, err)
+			paths, err := query.CollectAll(seq)
+			require.NoError(b, err)
+			require.Len(b, paths, 1)
+			require.Equal(b, "slow", paths[0].Subject.ObjectID)
+		}
+	})
 }