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

## Description

Currently, there are very few calls in pkg/query direct to the datastore interface. We'd like to keep this interface minimal, in an effort to start to clean up datastores over time as we continue with query planner.

This has the added benefit of being a way to inject various testing scenarios directly with the query planner tests and benchmarks; for example, adding a slight delay on any call to simulate a network latency.

Author: barakmich

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 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
 	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)
+		}
+	})
 }