Skip loading of caveats in SQL when unnecessary

internal/datastore/common/relationships.go

@@ -78,11 +78,19 @@ func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInf
 	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetObjectID, &subjectObjectID)
 	colsToSelect = StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetRelation, &subjectRelation)
 
-	colsToSelect = append(colsToSelect, &caveatName, &caveatCtx)
+	if !queryInfo.SkipCaveats {
+		colsToSelect = append(colsToSelect, &caveatName, &caveatCtx)
+	}
+
 	if withIntegrity {
 		colsToSelect = append(colsToSelect, &integrityKeyID, &integrityHash, &timestamp)
 	}
 
+	if len(colsToSelect) == 0 {
+		var unused int
+		colsToSelect = append(colsToSelect, &unused)
+	}
+
 	return func(yield func(tuple.Relationship, error) bool) {
 		err := tx.QueryFunc(ctx, func(ctx context.Context, rows R) error {
 			var r Rows = rows
@@ -100,11 +108,13 @@ func QueryRelationships[R Rows, C ~map[string]any](ctx context.Context, queryInf
 				}
 
 				var caveat *corev1.ContextualizedCaveat
-				if caveatName.Valid {
-					var err error
-					caveat, err = ContextualizedCaveatFrom(caveatName.String, caveatCtx)
-					if err != nil {
-						return fmt.Errorf(errUnableToQueryRels, fmt.Errorf("unable to fetch caveat context: %w", err))
+				if !queryInfo.SkipCaveats {
+					if caveatName.Valid {
+						var err error
+						caveat, err = ContextualizedCaveatFrom(caveatName.String, caveatCtx)
+						if err != nil {
+							return fmt.Errorf(errUnableToQueryRels, fmt.Errorf("unable to fetch caveat context: %w", err))
+						}
 					}
 				}
 

internal/datastore/common/sql.go

@@ -670,8 +670,16 @@ func (tqs QueryExecutor) ExecuteQuery(
 	columnNamesToSelect = checkColumn(columnNamesToSelect, query.filteringColumnTracker, query.schema.ColUsersetObjectID)
 	columnNamesToSelect = checkColumn(columnNamesToSelect, query.filteringColumnTracker, query.schema.ColUsersetRelation)
 
-	columnNamesToSelect = append(columnNamesToSelect, query.schema.ColCaveatName, query.schema.ColCaveatContext)
+	if !queryOpts.SkipCaveats {
+		columnNamesToSelect = append(columnNamesToSelect, query.schema.ColCaveatName, query.schema.ColCaveatContext)
+	}
+
+	selectingNoColumns := false
 	columnNamesToSelect = append(columnNamesToSelect, query.schema.ExtraFields...)
+	if len(columnNamesToSelect) == 0 {
+		columnNamesToSelect = append(columnNamesToSelect, "1")
+		selectingNoColumns = true
+	}
 
 	toExecute.queryBuilder = toExecute.queryBuilder.Columns(columnNamesToSelect...)
 
@@ -687,7 +695,7 @@ func (tqs QueryExecutor) ExecuteQuery(
 		return nil, err
 	}
 
-	return tqs.Executor(ctx, QueryInfo{query.schema, query.filteringColumnTracker}, sql, args)
+	return tqs.Executor(ctx, QueryInfo{query.schema, query.filteringColumnTracker, queryOpts.SkipCaveats, selectingNoColumns}, sql, args)
 }
 
 func checkColumn(columns []string, tracker map[string]ColumnTracker, colName string) []string {
@@ -699,8 +707,10 @@ func checkColumn(columns []string, tracker map[string]ColumnTracker, colName str
 
 // QueryInfo holds the schema information and filtering values for a query.
 type QueryInfo struct {
-	Schema          SchemaInformation
-	FilteringValues map[string]ColumnTracker
+	Schema             SchemaInformation
+	FilteringValues    map[string]ColumnTracker
+	SkipCaveats        bool
+	SelectingNoColumns bool
 }
 
 // ExecuteReadRelsQueryFunc is a function that can be used to execute a single rendered SQL query.

internal/datastore/memdb/readonly.go

@@ -148,11 +148,11 @@ func (r *memdbReader) QueryRelationships(
 		fallthrough
 
 	case options.ByResource:
-		iter := newMemdbTupleIterator(filteredIterator, queryOpts.Limit)
+		iter := newMemdbTupleIterator(filteredIterator, queryOpts.Limit, queryOpts.SkipCaveats)
 		return iter, nil
 
 	case options.BySubject:
-		return newSubjectSortedIterator(filteredIterator, queryOpts.Limit)
+		return newSubjectSortedIterator(filteredIterator, queryOpts.Limit, queryOpts.SkipCaveats)
 
 	default:
 		return nil, spiceerrors.MustBugf("unsupported sort order: %v", queryOpts.Sort)
@@ -210,11 +210,11 @@ func (r *memdbReader) ReverseQueryRelationships(
 		fallthrough
 
 	case options.ByResource:
-		iter := newMemdbTupleIterator(filteredIterator, queryOpts.LimitForReverse)
+		iter := newMemdbTupleIterator(filteredIterator, queryOpts.LimitForReverse, false)
 		return iter, nil
 
 	case options.BySubject:
-		return newSubjectSortedIterator(filteredIterator, queryOpts.LimitForReverse)
+		return newSubjectSortedIterator(filteredIterator, queryOpts.LimitForReverse, false)
 
 	default:
 		return nil, spiceerrors.MustBugf("unsupported sort order: %v", queryOpts.SortForReverse)
@@ -467,7 +467,7 @@ func makeCursorFilterFn(after options.Cursor, order options.SortOrder) func(tpl
 	return noopCursorFilter
 }
 
-func newSubjectSortedIterator(it memdb.ResultIterator, limit *uint64) (datastore.RelationshipIterator, error) {
+func newSubjectSortedIterator(it memdb.ResultIterator, limit *uint64, skipCaveats bool) (datastore.RelationshipIterator, error) {
 	results := make([]tuple.Relationship, 0)
 
 	// Coalesce all of the results into memory
@@ -477,6 +477,10 @@ func newSubjectSortedIterator(it memdb.ResultIterator, limit *uint64) (datastore
 			return nil, err
 		}
 
+		if skipCaveats && rt.OptionalCaveat != nil {
+			return nil, spiceerrors.MustBugf("unexpected caveat in result for relationship: %v", rt)
+		}
+
 		results = append(results, rt)
 	}
 
@@ -513,7 +517,7 @@ func eq(lhsNamespace, lhsObjectID, lhsRelation string, rhs tuple.ObjectAndRelati
 	return lhsNamespace == rhs.ObjectType && lhsObjectID == rhs.ObjectID && lhsRelation == rhs.Relation
 }
 
-func newMemdbTupleIterator(it memdb.ResultIterator, limit *uint64) datastore.RelationshipIterator {
+func newMemdbTupleIterator(it memdb.ResultIterator, limit *uint64, skipCaveats bool) datastore.RelationshipIterator {
 	var count uint64
 	return func(yield func(tuple.Relationship, error) bool) {
 		for {
@@ -527,6 +531,11 @@ func newMemdbTupleIterator(it memdb.ResultIterator, limit *uint64) datastore.Rel
 			}
 
 			rt, err := foundRaw.(*relationship).Relationship()
+			if skipCaveats && rt.OptionalCaveat != nil {
+				yield(rt, fmt.Errorf("unexpected caveat in result for relationship: %v", rt))
+				return
+			}
+
 			if !yield(rt, err) {
 				return
 			}

internal/datastore/mysql/readwrite.go

@@ -16,6 +16,7 @@ import (
 	"github.com/ccoveille/go-safecast"
 	"github.com/go-sql-driver/mysql"
 	"github.com/jzelinskie/stringz"
+	"golang.org/x/exp/maps"
 
 	"github.com/authzed/spicedb/internal/datastore/common"
 	"github.com/authzed/spicedb/internal/datastore/revisions"
@@ -58,6 +59,8 @@ func (cc *structpbWrapper) Scan(val any) error {
 	if !ok {
 		return fmt.Errorf("unsupported type: %T", v)
 	}
+
+	maps.Clear(*cc)
 	return json.Unmarshal(v, &cc)
 }
 

internal/datastore/spanner/reader.go

@@ -204,7 +204,14 @@ func queryExecutor(txSource txFactory) common.ExecuteReadRelsQueryFunc {
 			colsToSelect = common.StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetObjectID, &subjectObjectID)
 			colsToSelect = common.StaticValueOrAddColumnForSelect(colsToSelect, queryInfo, queryInfo.Schema.ColUsersetRelation, &subjectRelation)
 
-			colsToSelect = append(colsToSelect, &caveatName, &caveatCtx)
+			if !queryInfo.SkipCaveats {
+				colsToSelect = append(colsToSelect, &caveatName, &caveatCtx)
+			}
+
+			if len(colsToSelect) == 0 {
+				var unused int64
+				colsToSelect = append(colsToSelect, &unused)
+			}
 
 			if err := iter.Do(func(row *spanner.Row) error {
 				err := row.Columns(colsToSelect...)

internal/dispatch/graph/check_test.go

@@ -1251,7 +1251,7 @@ func TestCheckPermissionOverSchema(t *testing.T) {
   definition user {}
 
   definition role {
-    relation member: user
+    relation member: user with somecaveat
   }
 
   definition resource {
@@ -1287,7 +1287,7 @@ func TestCheckPermissionOverSchema(t *testing.T) {
   definition user {}
   
   definition role {
-    relation member: user
+    relation member: user with somecaveat
   }
 
   definition resource {

internal/graph/check.go

@@ -18,6 +18,7 @@ import (
 	"github.com/authzed/spicedb/internal/namespace"
 	"github.com/authzed/spicedb/internal/taskrunner"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	nspkg "github.com/authzed/spicedb/pkg/namespace"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -310,6 +311,8 @@ func (cc *ConcurrentChecker) checkDirect(ctx context.Context, crc currentRequest
 	hasNonTerminals := false
 	hasDirectSubject := false
 	hasWildcardSubject := false
+	directSubjectOrWildcardCanHaveCaveats := false
+	nonTerminalsCanHaveCaveats := false
 
 	defer func() {
 		if hasNonTerminals {
@@ -334,6 +337,10 @@ func (cc *ConcurrentChecker) checkDirect(ctx context.Context, crc currentRequest
 			} else if allowedDirectRelation.GetRelation() == crc.parentReq.Subject.Relation {
 				hasDirectSubject = true
 			}
+
+			if allowedDirectRelation.RequiredCaveat != nil {
+				directSubjectOrWildcardCanHaveCaveats = true
+			}
 		}
 
 		// If the relation found is not an ellipsis, then this is a nested relation that
@@ -343,6 +350,9 @@ func (cc *ConcurrentChecker) checkDirect(ctx context.Context, crc currentRequest
 		// relations can reach the target subject type.
 		if allowedDirectRelation.GetRelation() != tuple.Ellipsis {
 			hasNonTerminals = true
+			if allowedDirectRelation.RequiredCaveat != nil {
+				nonTerminalsCanHaveCaveats = true
+			}
 		}
 	}
 
@@ -381,7 +391,7 @@ func (cc *ConcurrentChecker) checkDirect(ctx context.Context, crc currentRequest
 			OptionalSubjectsSelectors: subjectSelectors,
 		}
 
-		it, err := ds.QueryRelationships(ctx, filter)
+		it, err := ds.QueryRelationships(ctx, filter, options.WithSkipCaveats(!directSubjectOrWildcardCanHaveCaveats))
 		if err != nil {
 			return checkResultError(NewCheckFailureErr(err), emptyMetadata)
 		}
@@ -430,7 +440,7 @@ func (cc *ConcurrentChecker) checkDirect(ctx context.Context, crc currentRequest
 		},
 	}
 
-	it, err := ds.QueryRelationships(ctx, filter)
+	it, err := ds.QueryRelationships(ctx, filter, options.WithSkipCaveats(!nonTerminalsCanHaveCaveats))
 	if err != nil {
 		return checkResultError(NewCheckFailureErr(err), emptyMetadata)
 	}

internal/testfixtures/datastore.go

@@ -37,6 +37,7 @@ var DocumentNS = ns.Namespace(
 	ns.MustRelation("owner",
 		nil,
 		ns.AllowedRelation("user", "..."),
+		ns.AllowedRelationWithCaveat("user", "...", ns.AllowedCaveat("test")),
 	),
 	ns.MustRelation("editor",
 		nil,
@@ -45,6 +46,7 @@ var DocumentNS = ns.Namespace(
 	ns.MustRelation("viewer",
 		nil,
 		ns.AllowedRelation("user", "..."),
+		ns.AllowedRelationWithCaveat("user", "...", ns.AllowedCaveat("test")),
 	),
 	ns.MustRelation("viewer_and_editor",
 		nil,
@@ -81,6 +83,7 @@ var FolderNS = ns.Namespace(
 	ns.MustRelation("owner",
 		nil,
 		ns.AllowedRelation("user", "..."),
+		ns.AllowedRelationWithCaveat("user", "...", ns.AllowedCaveat("test")),
 	),
 	ns.MustRelation("editor",
 		nil,
@@ -90,6 +93,7 @@ var FolderNS = ns.Namespace(
 		nil,
 		ns.AllowedRelation("user", "..."),
 		ns.AllowedRelation("folder", "viewer"),
+		ns.AllowedRelationWithCaveat("folder", "viewer", ns.AllowedCaveat("test")),
 	),
 	ns.MustRelation("parent", nil, ns.AllowedRelation("folder", "...")),
 	ns.MustRelation("edit",

pkg/datastore/options/options.go

@@ -43,9 +43,10 @@ func ToRelationship(c Cursor) *tuple.Relationship {
 
 // QueryOptions are the options that can affect the results of a normal forward query.
 type QueryOptions struct {
-	Limit *uint64   `debugmap:"visible"`
-	Sort  SortOrder `debugmap:"visible"`
-	After Cursor    `debugmap:"visible"`
+	Limit       *uint64   `debugmap:"visible"`
+	Sort        SortOrder `debugmap:"visible"`
+	After       Cursor    `debugmap:"visible"`
+	SkipCaveats bool      `debugmap:"visible"`
 }
 
 // ReverseQueryOptions are the options that can affect the results of a reverse query.

pkg/datastore/test/relationships.go

@@ -1003,10 +1003,11 @@ func RecreateRelationshipsAfterDeleteWithFilter(t *testing.T, tester DatastoreTe
 // QueryRelationshipsWithVariousFiltersTest tests various relationship filters for query relationships.
 func QueryRelationshipsWithVariousFiltersTest(t *testing.T, tester DatastoreTester) {
 	tcs := []struct {
-		name          string
-		filter        datastore.RelationshipsFilter
-		relationships []string
-		expected      []string
+		name           string
+		filter         datastore.RelationshipsFilter
+		withoutCaveats bool
+		relationships  []string
+		expected       []string
 	}{
 		{
 			name: "resource type",
@@ -1352,6 +1353,37 @@ func QueryRelationshipsWithVariousFiltersTest(t *testing.T, tester DatastoreTest
 				"folder:someotherfolder#viewer@user:tom",
 			},
 		},
+		{
+			name: "resource type with caveats",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType: "document",
+			},
+			relationships: []string{
+				"document:first#viewer@user:tom[firstcaveat]",
+				"document:second#viewer@user:tom[secondcaveat]",
+				"folder:secondfolder#viewer@user:tom",
+				"folder:someotherfolder#viewer@user:tom",
+			},
+			expected: []string{"document:first#viewer@user:tom[firstcaveat]", "document:second#viewer@user:tom[secondcaveat]"},
+		},
+		{
+			name: "resource type with caveats and context",
+			filter: datastore.RelationshipsFilter{
+				OptionalResourceType: "document",
+			},
+			relationships: []string{
+				"document:first#viewer@user:tom[firstcaveat:{\"foo\":\"bar\"}]",
+				"document:second#viewer@user:tom[secondcaveat]",
+				"document:third#viewer@user:tom[secondcaveat:{\"bar\":\"baz\"}]",
+				"folder:secondfolder#viewer@user:tom",
+				"folder:someotherfolder#viewer@user:tom",
+			},
+			expected: []string{
+				"document:first#viewer@user:tom[firstcaveat:{\"foo\":\"bar\"}]",
+				"document:second#viewer@user:tom[secondcaveat]",
+				"document:third#viewer@user:tom[secondcaveat:{\"bar\":\"baz\"}]",
+			},
+		},
 	}
 
 	for _, tc := range tcs {
@@ -1375,7 +1407,7 @@ func QueryRelationshipsWithVariousFiltersTest(t *testing.T, tester DatastoreTest
 			require.NoError(err)
 
 			reader := ds.SnapshotReader(headRev)
-			iter, err := reader.QueryRelationships(ctx, tc.filter)
+			iter, err := reader.QueryRelationships(ctx, tc.filter, options.WithSkipCaveats(tc.withoutCaveats))
 			require.NoError(err)
 
 			var results []string