chore(services): store observed counts in permissionsServer, add LR/LS flags (#2929)

Author: barakmich

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - 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)
 - Query Planner optimization: in Check requests, prune branches that cannot lead to the subject type specified (https://github.com/authzed/spicedb/pull/2968)
+- Added `lr` and `ls` to `--experimental-query-plan` for those endpoints, as well as in-memory statistics for optimizing the plans (https://github.com/authzed/spicedb/pull/2929)
 
 ### Fixed
 - Regression introduced in 1.49.2: missing spans in ReadSchema calls (https://github.com/authzed/spicedb/pull/2947)

docs/spicedb.md

@@ -234,6 +234,180 @@ func TestConsistency(t *testing.T) {
 	})
 }
 
+// TestQueryPlanConsistencyGRPC verifies that the experimental query plan engine
+// produces consistent results for LookupResources and LookupSubjects when those
+// operations are routed through the permissionServer gRPC handlers
+// (lookupResourcesWithQueryPlan / lookupSubjectsWithQueryPlan).
+//
+// It iterates the same validation files as TestConsistency, spins up a server
+// with both "lr" and "ls" query-plan flags enabled, and then verifies that the
+// set of returned resource/subject IDs matches the expected accessibility set.
+//
+// Note: the query plan handlers do not yet support cursor pagination or wildcard
+// exclusions, so this test intentionally uses unpaginated single-shot calls and
+// only checks set membership rather than running the full validateLookupResources
+// / validateLookupSubjects suite.
+func TestQueryPlanConsistencyGRPC(t *testing.T) { // nolint:tparallel
+	t.Parallel()
+
+	consistencyTestFiles, err := consistencytestutil.ListTestConfigs()
+	require.NoError(t, err)
+
+	for _, filePath := range consistencyTestFiles {
+		t.Run(path.Base(filePath), func(t *testing.T) {
+			t.Parallel()
+			runQueryPlanConsistencyGRPCForFile(t, filePath)
+		})
+	}
+}
+
+// runQueryPlanConsistencyGRPCForFile spins up a full gRPC server with both the
+// "lr" and "ls" experimental query plan flags enabled, then validates that
+// LookupResources and LookupSubjects return the correct set of IDs for every
+// resource-type × subject combination in the validation file.
+func runQueryPlanConsistencyGRPCForFile(t *testing.T, filePath string) {
+	options := []server.ConfigOption{
+		server.WithDispatchChunkSize(10),
+		server.WithExperimentalQueryPlan("lr"),
+		server.WithExperimentalQueryPlan("ls"),
+	}
+
+	cad := consistencytestutil.LoadDataAndCreateClusterForTesting(t, filePath, testTimedelta, options...)
+
+	headRevision, err := cad.DataStore.HeadRevision(cad.Ctx)
+	require.NoError(t, err)
+
+	accessibilitySet := consistencytestutil.BuildAccessibilitySet(t, cad.Ctx, cad.Populated, cad.DataStore)
+
+	testers := consistencytestutil.ServiceTesters(cad.Conn)
+	for _, tester := range testers {
+		t.Run(tester.Name(), func(t *testing.T) {
+			validateQueryPlanLookupResources(t, cad, tester, headRevision, accessibilitySet)
+			validateQueryPlanLookupSubjects(t, cad, tester, headRevision, accessibilitySet)
+		})
+	}
+}
+
+// validateQueryPlanLookupResources checks that for each resource-type × subject
+// pair the query plan handler returns exactly the same resource IDs as the
+// accessibility set expects. Pagination is not tested here because the query
+// plan handlers do not yet implement cursor-based pagination.
+func validateQueryPlanLookupResources(
+	t *testing.T,
+	cad consistencytestutil.ConsistencyClusterAndData,
+	tester consistencytestutil.ServiceTester,
+	revision datastore.Revision,
+	accessibilitySet *consistencytestutil.AccessibilitySet,
+) {
+	testForEachResourceTypeInPopulated(t, cad.Populated, "qp_lookup_resources",
+		func(t *testing.T, resourceRelation tuple.RelationReference) {
+			for _, subject := range accessibilitySet.AllSubjectsNoWildcards() {
+				t.Run(tuple.StringONR(subject), func(t *testing.T) {
+					accessibleResources := accessibilitySet.LookupAccessibleResources(resourceRelation, subject)
+
+					// Single unpaginated call (limit=0 means no limit).
+					foundResources, _, err := tester.LookupResources(t.Context(), resourceRelation, subject, revision, nil, 0, nil)
+					require.NoError(t, err)
+
+					resolvedIDs := make([]string, 0, len(foundResources))
+					for _, r := range foundResources {
+						resolvedIDs = append(resolvedIDs, r.ResourceObjectId)
+					}
+
+					require.ElementsMatch(t,
+						slices.Collect(maps.Keys(accessibleResources)),
+						resolvedIDs,
+						"query plan LookupResources mismatch for %s#%s / subject %s",
+						resourceRelation.ObjectType, resourceRelation.Relation, tuple.StringONR(subject),
+					)
+				})
+			}
+		})
+}
+
+// validateQueryPlanLookupSubjects checks that for each resource × subject-type
+// pair the query plan handler returns at least the subjects the accessibility
+// set defines as directly accessible. Wildcard exclusion checking is omitted
+// because the query plan handler does not yet populate ExcludedSubjects.
+func validateQueryPlanLookupSubjects(
+	t *testing.T,
+	cad consistencytestutil.ConsistencyClusterAndData,
+	tester consistencytestutil.ServiceTester,
+	revision datastore.Revision,
+	accessibilitySet *consistencytestutil.AccessibilitySet,
+) {
+	testForEachResourceInPopulated(t, cad.Populated, accessibilitySet, "qp_lookup_subjects",
+		func(t *testing.T, resource tuple.ObjectAndRelation) {
+			for _, subjectType := range accessibilitySet.SubjectTypes() {
+				t.Run(fmt.Sprintf("%s#%s", subjectType.ObjectType, subjectType.Relation),
+					func(t *testing.T) {
+						resolvedSubjects, err := tester.LookupSubjects(t.Context(), resource, subjectType, revision, nil)
+						require.NoError(t, err)
+
+						// The accessibility set only covers directly-accessible defined subjects;
+						// it does not include inferred subjects or wildcards, so we check subset.
+						expectedDefinedSubjects := accessibilitySet.DirectlyAccessibleDefinedSubjectsOfType(resource, subjectType)
+						requireSubsetOf(t,
+							slices.Collect(maps.Keys(resolvedSubjects)),
+							slices.Collect(maps.Keys(expectedDefinedSubjects)),
+						)
+					})
+			}
+		})
+}
+
+// testForEachResourceTypeInPopulated runs a subtest for every relation on every
+// namespace defined in the populated validation file.
+func testForEachResourceTypeInPopulated(
+	t *testing.T,
+	populated *validationfile.PopulatedValidationFile,
+	prefix string,
+	handler func(t *testing.T, resourceRelation tuple.RelationReference),
+) {
+	t.Helper()
+	for _, resourceType := range populated.NamespaceDefinitions {
+		for _, relation := range resourceType.Relation {
+			t.Run(fmt.Sprintf("%s_%s_%s", prefix, resourceType.Name, relation.Name),
+				func(t *testing.T) {
+					handler(t, tuple.RelationReference{
+						ObjectType: resourceType.Name,
+						Relation:   relation.Name,
+					})
+				})
+		}
+	}
+}
+
+// testForEachResourceInPopulated runs a subtest for every resource instance
+// present in the accessibility set, filtered by namespace definitions.
+func testForEachResourceInPopulated(
+	t *testing.T,
+	populated *validationfile.PopulatedValidationFile,
+	accessibilitySet *consistencytestutil.AccessibilitySet,
+	prefix string,
+	handler func(t *testing.T, resource tuple.ObjectAndRelation),
+) {
+	t.Helper()
+	for _, resourceType := range populated.NamespaceDefinitions {
+		resources, ok := accessibilitySet.ResourcesByNamespace.Get(resourceType.Name)
+		if !ok {
+			continue
+		}
+		for _, relation := range resourceType.Relation {
+			for _, resource := range resources {
+				t.Run(fmt.Sprintf("%s_%s_%s_%s", prefix, resourceType.Name, resource.ObjectID, relation.Name),
+					func(t *testing.T) {
+						handler(t, tuple.ObjectAndRelation{
+							ObjectType: resourceType.Name,
+							ObjectID:   resource.ObjectID,
+							Relation:   relation.Name,
+						})
+					})
+			}
+		}
+	}
+}
+
 func runConsistencyTestSuiteForFile(t *testing.T, filePath string, useCachingDispatcher bool, chunkSize uint16) {
 	options := []server.ConfigOption{
 		server.WithDispatchChunkSize(chunkSize),

internal/services/integrationtesting/consistency_test.go

@@ -78,20 +78,23 @@ func runQueryPlanConsistencyForFile(t *testing.T, filePath string) {
 
 	t.Run("lookup_resources", func(t *testing.T) {
 		if os.Getenv("TEST_QUERY_PLAN_RESOURCES") == "" {
-			t.Skip("Skipping IterResources tests: set TEST_QUERY_PLAN_RESOURCES=true to enable")
+			t.Skip("Skipping IterResources tests due to deprectation: Set TEST_QUERY_PLAN_RESOURCES=true to enable")
 		}
 		runQueryPlanLookupResources(t, handle)
 	})
 
 	t.Run("lookup_subjects", func(t *testing.T) {
 		if os.Getenv("TEST_QUERY_PLAN_SUBJECTS") == "" {
-			t.Skip("Skipping IterSubjects tests: set TEST_QUERY_PLAN_SUBJECTS=true to enable")
+			t.Skip("Skipping IterSubjects tests due to deprectation: Set TEST_QUERY_PLAN_SUBJECTS=true to enable")
 		}
 		runQueryPlanLookupSubjects(t, handle)
 	})
 }
 
 func runQueryPlanAssertions(t *testing.T, handle *queryPlanConsistencyHandle) {
+	if os.Getenv("TEST_QUERY_PLAN_CHECK") == "" {
+		t.Skip("Skipping Check tests due to deprecation. Set TEST_QUERY_PLAN_CHECK=true to enable")
+	}
 	t.Run("assertions", func(t *testing.T) {
 		for _, parsedFile := range handle.populated.ParsedFiles {
 			for _, entry := range []struct {

internal/services/integrationtesting/query_plan_consistency_test.go

@@ -70,7 +70,7 @@ func (ps *permissionServer) CheckPermission(ctx context.Context, req *v1.CheckPe
 
 	telemetry.LogicalChecks.Inc()
 
-	if ps.config.ExperimentalQueryPlan {
+	if ps.config.ExperimentalQueryPlan.Check {
 		return ps.checkPermissionWithQueryPlan(ctx, req)
 	}
 
@@ -475,6 +475,10 @@ func (ps *permissionServer) LookupResources(req *v1.LookupResourcesRequest, resp
 		}
 	}
 
+	if ps.config.ExperimentalQueryPlan.LookupResources {
+		return ps.lookupResourcesWithQueryPlan(req, resp)
+	}
+
 	if ps.config.EnableExperimentalLookupResources3 {
 		return ps.lookupResources3(req, resp)
 	}
@@ -792,6 +796,10 @@ func (ps *permissionServer) LookupSubjects(req *v1.LookupSubjectsRequest, resp v
 
 	ctx := resp.Context()
 
+	if ps.config.ExperimentalQueryPlan.LookupSubjects {
+		return ps.lookupSubjectsWithQueryPlan(req, resp)
+	}
+
 	if req.OptionalConcreteLimit != 0 {
 		return ps.rewriteError(ctx, status.Errorf(codes.Unimplemented, "concrete limit is not yet supported"))
 	}