Skip to content

chore: add ctx to RevisionedReader.ReadSchema() #2952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
miparnisari merged 1 commit into from
Mar 6, 2026
Merged

chore: add ctx to RevisionedReader.ReadSchema() #2952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions internal/caveats/run_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -471,7 +471,7 @@ func TestRunCaveatExpressions(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

for _, debugOption := range []RunCaveatExpressionDebugOption{
Expand Down Expand Up @@ -524,7 +524,7 @@ func TestRunCaveatWithMissingMap(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

result, err := RunSingleCaveatExpression(
Expand Down Expand Up @@ -556,7 +556,7 @@ func TestRunCaveatWithEmptyMap(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

_, err = RunSingleCaveatExpression(
Expand Down Expand Up @@ -594,7 +594,7 @@ func TestRunCaveatMultipleTimes(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

runner := NewCaveatRunner(types.Default.TypeSet)
Expand Down Expand Up @@ -662,7 +662,7 @@ func TestRunCaveatWithMissingDefinition(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

// Try to run a caveat that doesn't exist
Expand Down Expand Up @@ -697,7 +697,7 @@ func TestCaveatRunnerPopulateCaveatDefinitionsForExpr(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

runner := NewCaveatRunner(types.Default.TypeSet)
Expand Down Expand Up @@ -742,7 +742,7 @@ func TestCaveatRunnerEmptyExpression(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

runner := NewCaveatRunner(types.Default.TypeSet)
Expand Down Expand Up @@ -823,7 +823,7 @@ func TestUnknownCaveatOperation(t *testing.T) {
req.NoError(err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(headRevision).ReadSchema()
sr, err := dl.SnapshotReader(headRevision).ReadSchema(t.Context())
req.NoError(err)

runner := NewCaveatRunner(types.Default.TypeSet)
Expand Down
2 changes: 1 addition & 1 deletion internal/dispatch/graph/graph.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ func (ld *localDispatcher) loadNamespace(ctx context.Context, nsName string, rev
reader := datalayer.MustFromContext(ctx).SnapshotReader(revision)

// Load namespace and relation from the datastore
schemaReader, err := reader.ReadSchema()
schemaReader, err := reader.ReadSchema(ctx)
if err != nil {
return nil, rewriteNamespaceError(err)
}
Expand Down
2 changes: 1 addition & 1 deletion internal/dispatch/keys/keys.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ func (c *CanonicalKeyHandler) CheckCacheKey(ctx context.Context, req *v1.Dispatc
}
r := dl.SnapshotReader(revision)

sr, err := r.ReadSchema()
sr, err := r.ReadSchema(ctx)
if err != nil {
return emptyDispatchCacheKey, err
}
Expand Down
4 changes: 2 additions & 2 deletions internal/graph/check.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -659,7 +659,7 @@ func (cc *ConcurrentChecker) checkComputedUserset(ctx context.Context, crc curre
// the same namespace as the caller, and thus must be fully typed checked.
if cu.Object == core.ComputedUserset_TUPLE_USERSET_OBJECT {
dl := datalayer.MustFromContext(ctx).SnapshotReader(crc.parentReq.Revision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return checkResultError(err, emptyMetadata)
}
Expand Down Expand Up @@ -698,7 +698,7 @@ type Traits struct {
// types of the given relation support caveats or expiration.
func TraitsForArrowRelation(ctx context.Context, reader datalayer.RevisionedReader, namespaceName string, relationName string) (Traits, error) {
// TODO(jschorr): Change to use the type system once we wire it through Check dispatch.
schemaReader, err := reader.ReadSchema()
schemaReader, err := reader.ReadSchema(ctx)
if err != nil {
return Traits{}, err
}
Expand Down
2 changes: 1 addition & 1 deletion internal/graph/computed/computecheck.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ func computeCaveatedCheckResult(ctx context.Context, runner *cexpr.CaveatRunner,

dl := datalayer.MustFromContext(ctx)
reader := dl.SnapshotReader(params.AtRevision)
sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return nil, err
}
Expand Down
2 changes: 1 addition & 1 deletion internal/graph/expand.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ func (ce *ConcurrentExpander) expandComputedUserset(ctx context.Context, req Val

// Check if the target relation exists. If not, return nothing.
dl := datalayer.MustFromContext(ctx).SnapshotReader(req.Revision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return expandError(err)
}
Expand Down
4 changes: 2 additions & 2 deletions internal/graph/lookupresources2.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ func (crr *CursoredLookupResources2) afterSameType(
// Load the type system and reachability graph to find the entrypoints for the reachability.
dl := datalayer.MustFromContext(ctx)
reader := dl.SnapshotReader(req.Revision)
sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -363,7 +363,7 @@ func (crr *CursoredLookupResources2) redispatchOrReportOverDatabaseQuery(
toBeHandled := make([]itemAndPostCursor[dispatchableResourcesSubjectMap2], 0)
currentCursor := queryCursor
caveatRunner := caveats.NewCaveatRunner(crr.caveatTypeSet)
caveatSR, err := config.reader.ReadSchema()
caveatSR, err := config.reader.ReadSchema(ctx)
if err != nil {
return nil, err
}
Expand Down
4 changes: 2 additions & 2 deletions internal/graph/lookupresources3.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ func (crr *CursoredLookupResources3) LookupResources3(req ValidatedLookupResourc
// interfaces used by various suboperations of the lookup resources operation.
dl := datalayer.MustFromContext(stream.Context())
reader := dl.SnapshotReader(req.Revision)
sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -813,7 +813,7 @@ func (crr *CursoredLookupResources3) relationshipsIter(

// Start the new relationships chunk that we will fill as we iterate over the results.
// It starts at the given dbCursor, which may be nil/empty if this is the first chunk.
caveatSR, err := refs.reader.ReadSchema()
caveatSR, err := refs.reader.ReadSchema(ctx)
if err != nil {
yieldError(err)
return
Expand Down
8 changes: 4 additions & 4 deletions internal/graph/lookupsubjects.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ func (cl *ConcurrentLookupSubjects) LookupSubjects(

dl := datalayer.MustFromContext(ctx)
reader := dl.SnapshotReader(req.Revision)
sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -199,7 +199,7 @@ func (cl *ConcurrentLookupSubjects) lookupViaComputed(
cu *core.ComputedUserset,
) error {
dl := datalayer.MustFromContext(ctx).SnapshotReader(parentRequest.Revision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -259,7 +259,7 @@ func lookupViaIntersectionTupleToUserset(
ttu *core.FunctionedTupleToUserset,
) error {
dl := datalayer.MustFromContext(ctx).SnapshotReader(parentRequest.Revision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -438,7 +438,7 @@ func lookupViaTupleToUserset[T relation](
relationshipsBySubjectONR := mapz.NewMultiMap[tuple.ObjectAndRelation, tuple.Relationship]()

dl := datalayer.MustFromContext(ctx).SnapshotReader(parentRequest.Revision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down
2 changes: 1 addition & 1 deletion internal/namespace/util_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ func TestCheckNamespaceAndRelations(t *testing.T) {
require.NoError(t, err)

dl := datalayer.NewDataLayer(ds)
sr, err := dl.SnapshotReader(rev).ReadSchema()
sr, err := dl.SnapshotReader(rev).ReadSchema(t.Context())
require.NoError(t, err)

err = namespace.CheckNamespaceAndRelations(t.Context(), tc.checks, sr)
Expand Down
2 changes: 1 addition & 1 deletion internal/relationships/validation_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,7 +336,7 @@ func TestValidateRelationshipOperations(t *testing.T) {

uds, rev := testfixtures.DatastoreFromSchemaAndTestRelationships(ds, tc.schema, nil, req)
dl := datalayer.NewDataLayer(uds)
sr, err := dl.SnapshotReader(rev).ReadSchema()
sr, err := dl.SnapshotReader(rev).ReadSchema(t.Context())
req.NoError(err)

op := tuple.Create
Expand Down
4 changes: 2 additions & 2 deletions internal/services/shared/schema.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ type AppliedSchemaChanges struct {
// ApplySchemaChanges applies schema changes found in the validated changes struct, via the specified
// ReadWriteTransaction.
func ApplySchemaChanges(ctx context.Context, rwt datalayer.ReadWriteTransaction, caveatTypeSet *caveattypes.TypeSet, validated *ValidatedSchemaChanges) (*AppliedSchemaChanges, error) {
sr, err := rwt.ReadSchema()
sr, err := rwt.ReadSchema(ctx)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -231,7 +231,7 @@ func ApplySchemaChangesOverExisting(

// Get the list of extant definitions so that we can add them to the
// list of definitions that should be written in the single shot
sr, err := rwt.ReadSchema()
sr, err := rwt.ReadSchema(ctx)
if err != nil {
return nil, err
}
Expand Down
2 changes: 1 addition & 1 deletion internal/services/shared/schema_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -776,7 +776,7 @@ func TestApplySchemaChangesOverExisting(t *testing.T) {
require.NoError(err)
require.Equal(tc.expectedAppliedSchemaChanges, *applied)

sr, err := rwt.ReadSchema()
sr, err := rwt.ReadSchema(ctx)
require.NoError(err)
schemaText, err := sr.SchemaText(t.Context())
require.NoError(err)
Expand Down
4 changes: 2 additions & 2 deletions internal/services/v1/bulkcheck.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ func (bc *bulkChecker) checkBulkPermissions(ctx context.Context, req *v1.CheckBu

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -253,7 +253,7 @@ func (bc *bulkChecker) checkBulkPermissions(ctx context.Context, req *v1.CheckBu

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return appendResultsForError(group.params, resourceIDs, err)
}
Expand Down
10 changes: 5 additions & 5 deletions internal/services/v1/experimental.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ func (es *experimentalServer) BulkImportRelationships(stream v1.ExperimentalServ
caveat: core.ContextualizedCaveat{},
caveatTypeSet: es.caveatTypeSet,
}
sr, err := rwt.ReadSchema()
sr, err := rwt.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -352,7 +352,7 @@ func BulkExport(ctx context.Context, dl datalayer.DataLayer, batchSize uint64, r

reader := dl.SnapshotReader(atRevision)

sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return shared.RewriteErrorWithoutConfig(ctx, err)
}
Expand Down Expand Up @@ -603,7 +603,7 @@ func (es *experimentalServer) ExperimentalComputablePermissions(ctx context.Cont
}

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return nil, shared.RewriteErrorWithoutConfig(ctx, err)
}
Expand Down Expand Up @@ -690,7 +690,7 @@ func (es *experimentalServer) ExperimentalDependentRelations(ctx context.Context
}

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)
sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return nil, shared.RewriteErrorWithoutConfig(ctx, err)
}
Expand Down Expand Up @@ -764,7 +764,7 @@ func (es *experimentalServer) ExperimentalRegisterRelationshipCounter(ctx contex
}

_, err := dl.ReadWriteTx(ctx, func(ctx context.Context, rwt datalayer.ReadWriteTransaction) error {
sr, err := rwt.ReadSchema()
sr, err := rwt.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down
14 changes: 7 additions & 7 deletions internal/services/v1/permissions.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ func (ps *permissionServer) CheckPermission(ctx context.Context, req *v1.CheckPe
return nil, ps.rewriteError(ctx, err)
}

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return nil, ps.rewriteError(ctx, err)
}
Expand Down Expand Up @@ -253,7 +253,7 @@ func (ps *permissionServer) ExpandPermissionTree(ctx context.Context, req *v1.Ex

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return nil, ps.rewriteError(ctx, err)
}
Expand Down Expand Up @@ -496,7 +496,7 @@ func (ps *permissionServer) lookupResources3(req *v1.LookupResourcesRequest, res

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return ps.rewriteError(ctx, err)
}
Expand Down Expand Up @@ -647,7 +647,7 @@ func (ps *permissionServer) lookupResources2(req *v1.LookupResourcesRequest, res

dl := datalayer.MustFromContext(ctx).SnapshotReader(atRevision)

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return ps.rewriteError(ctx, err)
}
Expand Down Expand Up @@ -808,7 +808,7 @@ func (ps *permissionServer) LookupSubjects(req *v1.LookupSubjectsRequest, resp v
return ps.rewriteError(ctx, err)
}

sr, err := dl.ReadSchema()
sr, err := dl.ReadSchema(ctx)
if err != nil {
return ps.rewriteError(ctx, err)
}
Expand Down Expand Up @@ -1096,7 +1096,7 @@ func (ps *permissionServer) ImportBulkRelationships(stream grpc.ClientStreamingS
numWritten += streamWritten

// The stream has terminated because we're awaiting namespace and/or caveat information
schemaReader, err := rwt.ReadSchema()
schemaReader, err := rwt.ReadSchema(ctx)
if err != nil {
return err
}
Expand Down Expand Up @@ -1184,7 +1184,7 @@ func ExportBulk(ctx context.Context, dl datalayer.DataLayer, batchSize uint64, r

reader := dl.SnapshotReader(atRevision)

readerSchema, err := reader.ReadSchema()
readerSchema, err := reader.ReadSchema(ctx)
if err != nil {
return shared.RewriteErrorWithoutConfig(ctx, err)
}
Expand Down
2 changes: 1 addition & 1 deletion internal/services/v1/permissions_queryplan.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@

// Load all namespace and caveat definitions to build the schema
// TODO: Better schema caching
sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)

Check warning on line 29 in internal/services/v1/permissions_queryplan.go

View check run for this annotation

Codecov / codecov/patch

internal/services/v1/permissions_queryplan.go#L29 

Added line #L29 was not covered by tests
if err != nil {
return nil, ps.rewriteError(ctx, err)
}
Expand Down
2 changes: 1 addition & 1 deletion internal/services/v1/reflectionutil.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ func loadCurrentSchema(ctx context.Context) (*diff.DiffableSchema, datastore.Rev

reader := dl.SnapshotReader(atRevision)

sr, err := reader.ReadSchema()
sr, err := reader.ReadSchema(ctx)
if err != nil {
return nil, atRevision, err
}
Expand Down
Loading
Loading