Add a concept of a caveat TypeSet to allow overriding the types available to caveat processing

internal/caveats/debug_test.go

@@ -251,7 +251,7 @@ func eval(expr string, context map[string]any) func(t *testing.T) ExpressionResu
 		env := caveats.NewEnvironment()
 
 		for name := range context {
-			require.NoError(t, env.AddVariable(name, types.AnyType))
+			require.NoError(t, env.AddVariable(name, types.Default.AnyType))
 		}
 
 		caveat, err := caveats.CompileCaveatWithName(env, expr, "test")

internal/caveats/run.go

@@ -35,24 +35,27 @@ const (
 // This instantiates its own CaveatRunner, and should therefore only be used in one-off situations.
 func RunSingleCaveatExpression(
 	ctx context.Context,
+	ts *caveattypes.TypeSet,
 	expr *core.CaveatExpression,
 	context map[string]any,
 	reader datastore.CaveatReader,
 	debugOption RunCaveatExpressionDebugOption,
 ) (ExpressionResult, error) {
-	runner := NewCaveatRunner()
+	runner := NewCaveatRunner(ts)
 	return runner.RunCaveatExpression(ctx, expr, context, reader, debugOption)
 }
 
 // CaveatRunner is a helper for running caveats, providing a cache for deserialized caveats.
 type CaveatRunner struct {
+	caveatTypeSet       *caveattypes.TypeSet
 	caveatDefs          map[string]*core.CaveatDefinition
 	deserializedCaveats map[string]*caveats.CompiledCaveat
 }
 
 // NewCaveatRunner creates a new CaveatRunner.
-func NewCaveatRunner() *CaveatRunner {
+func NewCaveatRunner(ts *caveattypes.TypeSet) *CaveatRunner {
 	return &CaveatRunner{
+		caveatTypeSet:       ts,
 		caveatDefs:          map[string]*core.CaveatDefinition{},
 		deserializedCaveats: map[string]*caveats.CompiledCaveat{},
 	}
@@ -130,7 +133,7 @@ func (cr *CaveatRunner) get(caveatDefName string) (*core.CaveatDefinition, *cave
 		return caveat, deserialized, nil
 	}
 
-	parameterTypes, err := caveattypes.DecodeParameterTypes(caveat.ParameterTypes)
+	parameterTypes, err := caveattypes.DecodeParameterTypes(cr.caveatTypeSet, caveat.ParameterTypes)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -185,6 +188,7 @@ func (cr *CaveatRunner) runExpressionWithCaveats(
 
 		// Perform type checking and conversion on the context map.
 		typedParameters, err := caveats.ConvertContextToParameters(
+			cr.caveatTypeSet,
 			untypedFullContext,
 			caveat.ParameterTypes,
 			caveats.SkipUnknownParameters,

internal/caveats/run_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/authzed/spicedb/internal/datastore/dsfortesting"
 	"github.com/authzed/spicedb/internal/datastore/memdb"
 	"github.com/authzed/spicedb/internal/testfixtures"
+	"github.com/authzed/spicedb/pkg/caveats/types"
 	"github.com/authzed/spicedb/pkg/datastore"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
@@ -478,7 +479,7 @@ func TestRunCaveatExpressions(t *testing.T) {
 				t.Run(fmt.Sprintf("%v", debugOption), func(t *testing.T) {
 					req := require.New(t)
 
-					result, err := caveats.RunSingleCaveatExpression(context.Background(), tc.expression, tc.context, reader, debugOption)
+					result, err := caveats.RunSingleCaveatExpression(context.Background(), types.Default.TypeSet, tc.expression, tc.context, reader, debugOption)
 					req.NoError(err)
 					req.Equal(tc.expectedValue, result.Value())
 
@@ -524,6 +525,7 @@ func TestRunCaveatWithMissingMap(t *testing.T) {
 
 	result, err := caveats.RunSingleCaveatExpression(
 		context.Background(),
+		types.Default.TypeSet,
 		caveatexpr("some_caveat"),
 		map[string]any{},
 		reader,
@@ -553,6 +555,7 @@ func TestRunCaveatWithEmptyMap(t *testing.T) {
 
 	_, err = caveats.RunSingleCaveatExpression(
 		context.Background(),
+		types.Default.TypeSet,
 		caveatexpr("some_caveat"),
 		map[string]any{
 			"themap": map[string]any{},
@@ -584,8 +587,7 @@ func TestRunCaveatMultipleTimes(t *testing.T) {
 	req.NoError(err)
 
 	reader := ds.SnapshotReader(headRevision)
-
-	runner := caveats.NewCaveatRunner()
+	runner := caveats.NewCaveatRunner(types.Default.TypeSet)
 
 	// Run the first caveat.
 	result, err := runner.RunCaveatExpression(context.Background(), caveatexpr("some_caveat"), map[string]any{

internal/datastore/proxy/schemacaching/estimatedsize_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/authzed/spicedb/internal/datastore/dsfortesting"
 	"github.com/authzed/spicedb/internal/datastore/memdb"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/validationfile"
 )
@@ -50,7 +51,7 @@ func TestEstimatedDefinitionSizes(t *testing.T) {
 			ds, err := dsfortesting.NewMemDBDatastoreForTesting(0, 1*time.Second, memdb.DisableGC)
 			require.NoError(err)
 
-			fullyResolved, _, err := validationfile.PopulateFromFiles(context.Background(), ds, []string{filePath})
+			fullyResolved, _, err := validationfile.PopulateFromFiles(context.Background(), ds, caveattypes.Default.TypeSet, []string{filePath})
 			require.NoError(err)
 
 			for _, nsDef := range fullyResolved.NamespaceDefinitions {

internal/datastore/proxy/schemacaching/standardcaching_test.go

@@ -357,7 +357,7 @@ func TestSnapshotCachingRealDatastore(t *testing.T) {
 			"document",
 			ns.MustCaveatDefinition(caveats.MustEnvForVariables(
 				map[string]caveattypes.VariableType{
-					"somevar": caveattypes.IntType,
+					"somevar": caveattypes.Default.IntType,
 				},
 			), "somecaveat", "somevar < 42"),
 			"somecaveat",

internal/dispatch/cluster/cluster.go

@@ -9,6 +9,7 @@ import (
 	"github.com/authzed/spicedb/internal/dispatch/keys"
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/cache"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 )
 
 // Option is a function-style option for configuring a combined Dispatcher.
@@ -21,6 +22,7 @@ type optionState struct {
 	concurrencyLimits     graph.ConcurrencyLimits
 	remoteDispatchTimeout time.Duration
 	dispatchChunkSize     uint16
+	caveatTypeSet         *caveattypes.TypeSet
 }
 
 // MetricsEnabled enables issuing prometheus metrics
@@ -66,6 +68,14 @@ func RemoteDispatchTimeout(remoteDispatchTimeout time.Duration) Option {
 	}
 }
 
+// CaveatTypeSet sets the type set to use for caveats. If not specified, the default
+// type set is used.
+func CaveatTypeSet(caveatTypeSet *caveattypes.TypeSet) Option {
+	return func(state *optionState) {
+		state.caveatTypeSet = caveatTypeSet
+	}
+}
+
 // NewClusterDispatcher takes a dispatcher (such as one created by
 // combined.NewDispatcher) and returns a cluster dispatcher suitable for use as
 // the dispatcher for the dispatch grpc server.
@@ -80,7 +90,9 @@ func NewClusterDispatcher(dispatch dispatch.Dispatcher, options ...Option) (disp
 		chunkSize = 100
 		log.Warn().Msgf("ClusterDispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
 	}
-	clusterDispatch := graph.NewDispatcher(dispatch, opts.concurrencyLimits, opts.dispatchChunkSize)
+
+	cts := caveattypes.TypeSetOrDefault(opts.caveatTypeSet)
+	clusterDispatch := graph.NewDispatcher(dispatch, cts, opts.concurrencyLimits, opts.dispatchChunkSize)
 
 	if opts.prometheusSubsystem == "" {
 		opts.prometheusSubsystem = "dispatch"

internal/dispatch/combined/combined.go

@@ -18,6 +18,7 @@ import (
 	"github.com/authzed/spicedb/internal/grpchelpers"
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/cache"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 )
 
@@ -38,6 +39,7 @@ type optionState struct {
 	secondaryUpstreamExprs      map[string]string
 	dispatchChunkSize           uint16
 	startingPrimaryHedgingDelay time.Duration
+	caveatTypeSet               *caveattypes.TypeSet
 }
 
 // MetricsEnabled enables issuing prometheus metrics
@@ -140,6 +142,14 @@ func StartingPrimaryHedgingDelay(startingPrimaryHedgingDelay time.Duration) Opti
 	}
 }
 
+// CaveatTypeSet sets the type set to use for caveats. If not specified, the default
+// type set is used.
+func CaveatTypeSet(caveatTypeSet *caveattypes.TypeSet) Option {
+	return func(state *optionState) {
+		state.caveatTypeSet = caveatTypeSet
+	}
+}
+
 // NewDispatcher initializes a Dispatcher that caches and redispatches
 // optionally to the provided upstream.
 func NewDispatcher(options ...Option) (dispatch.Dispatcher, error) {
@@ -163,7 +173,9 @@ func NewDispatcher(options ...Option) (dispatch.Dispatcher, error) {
 		chunkSize = 100
 		log.Warn().Msgf("CombinedDispatcher: dispatchChunkSize not set, defaulting to %d", chunkSize)
 	}
-	redispatch := graph.NewDispatcher(cachingRedispatch, opts.concurrencyLimits, chunkSize)
+
+	cts := caveattypes.TypeSetOrDefault(opts.caveatTypeSet)
+	redispatch := graph.NewDispatcher(cachingRedispatch, cts, opts.concurrencyLimits, chunkSize)
 	redispatch = singleflight.New(redispatch, &keys.CanonicalKeyHandler{})
 
 	// If an upstream is specified, create a cluster dispatcher.

internal/dispatch/graph/check_test.go

@@ -19,6 +19,7 @@ import (
 	log "github.com/authzed/spicedb/internal/logging"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	"github.com/authzed/spicedb/internal/testfixtures"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -170,7 +171,7 @@ func TestMaxDepth(t *testing.T) {
 	revision, err := common.UpdateRelationshipsInDatastore(ctx, ds, mutation)
 	require.NoError(err)
 
-	dispatch := NewLocalOnlyDispatcher(10, 100)
+	dispatch := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 	_, err = dispatch.DispatchCheck(ctx, &v1.DispatchCheckRequest{
 		ResourceRelation: RR("folder", "view").ToCoreRR(),
@@ -1328,7 +1329,7 @@ func TestCheckPermissionOverSchema(t *testing.T) {
 
 			require := require.New(t)
 
-			dispatcher := NewLocalOnlyDispatcher(10, 100)
+			dispatcher := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 			ds, err := dsfortesting.NewMemDBDatastoreForTesting(0, 0, memdb.DisableGC)
 			require.NoError(err)
@@ -1833,7 +1834,7 @@ func TestCheckWithHints(t *testing.T) {
 
 			require := require.New(t)
 
-			dispatcher := NewLocalOnlyDispatcher(10, 100)
+			dispatcher := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 			ds, err := dsfortesting.NewMemDBDatastoreForTesting(0, 0, memdb.DisableGC)
 			require.NoError(err)
@@ -1873,7 +1874,7 @@ func TestCheckHintsPartialApplication(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)
 
-	dispatcher := NewLocalOnlyDispatcher(10, 100)
+	dispatcher := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 	ds, err := dsfortesting.NewMemDBDatastoreForTesting(0, 0, memdb.DisableGC)
 	require.NoError(err)
@@ -1919,7 +1920,7 @@ func TestCheckHintsPartialApplicationOverArrow(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)
 
-	dispatcher := NewLocalOnlyDispatcher(10, 100)
+	dispatcher := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 	ds, err := dsfortesting.NewMemDBDatastoreForTesting(0, 0, memdb.DisableGC)
 	require.NoError(err)
@@ -1972,7 +1973,7 @@ func newLocalDispatcherWithConcurrencyLimit(t testing.TB, concurrencyLimit uint1
 
 	ds, revision := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 
-	dispatch := NewLocalOnlyDispatcher(concurrencyLimit, 100)
+	dispatch := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, concurrencyLimit, 100)
 
 	cachingDispatcher, err := caching.NewCachingDispatcher(caching.DispatchTestCache(t), false, "", &keys.CanonicalKeyHandler{})
 	require.NoError(t, err)
@@ -1994,7 +1995,7 @@ func newLocalDispatcherWithSchemaAndRels(t testing.TB, schema string, rels []tup
 
 	ds, revision := testfixtures.DatastoreFromSchemaAndTestRelationships(rawDS, schema, rels, require.New(t))
 
-	dispatch := NewLocalOnlyDispatcher(10, 100)
+	dispatch := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 	cachingDispatcher, err := caching.NewCachingDispatcher(caching.DispatchTestCache(t), false, "", &keys.CanonicalKeyHandler{})
 	require.NoError(t, err)

internal/dispatch/graph/expand_test.go

@@ -20,6 +20,7 @@ import (
 	expand "github.com/authzed/spicedb/internal/graph"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	"github.com/authzed/spicedb/internal/testfixtures"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	"github.com/authzed/spicedb/pkg/graph"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
@@ -295,7 +296,7 @@ func TestMaxDepthExpand(t *testing.T) {
 	require.NoError(err)
 	require.NoError(datastoremw.SetInContext(ctx, ds))
 
-	dispatch := NewLocalOnlyDispatcher(10, 100)
+	dispatch := NewLocalOnlyDispatcher(caveattypes.Default.TypeSet, 10, 100)
 
 	_, err = dispatch.DispatchExpand(ctx, &v1.DispatchExpandRequest{
 		ResourceAndRelation: tuple.CoreONR("folder", "oops", "view"),

internal/dispatch/graph/graph.go

@@ -16,6 +16,7 @@ import (
 	"github.com/authzed/spicedb/internal/graph"
 	log "github.com/authzed/spicedb/internal/logging"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/middleware/nodeid"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -79,13 +80,13 @@ func SharedConcurrencyLimits(concurrencyLimit uint16) ConcurrencyLimits {
 }
 
 // NewLocalOnlyDispatcher creates a dispatcher that consults with the graph to formulate a response.
-func NewLocalOnlyDispatcher(concurrencyLimit uint16, dispatchChunkSize uint16) dispatch.Dispatcher {
-	return NewLocalOnlyDispatcherWithLimits(SharedConcurrencyLimits(concurrencyLimit), dispatchChunkSize)
+func NewLocalOnlyDispatcher(typeSet *caveattypes.TypeSet, concurrencyLimit uint16, dispatchChunkSize uint16) dispatch.Dispatcher {
+	return NewLocalOnlyDispatcherWithLimits(typeSet, SharedConcurrencyLimits(concurrencyLimit), dispatchChunkSize)
 }
 
 // NewLocalOnlyDispatcherWithLimits creates a dispatcher thatg consults with the graph to formulate a response
 // and has the defined concurrency limits per dispatch type.
-func NewLocalOnlyDispatcherWithLimits(concurrencyLimits ConcurrencyLimits, dispatchChunkSize uint16) dispatch.Dispatcher {
+func NewLocalOnlyDispatcherWithLimits(typeSet *caveattypes.TypeSet, concurrencyLimits ConcurrencyLimits, dispatchChunkSize uint16) dispatch.Dispatcher {
 	d := &localDispatcher{}
 
 	concurrencyLimits = limitsOrDefaults(concurrencyLimits, defaultConcurrencyLimit)
@@ -98,14 +99,14 @@ func NewLocalOnlyDispatcherWithLimits(concurrencyLimits ConcurrencyLimits, dispa
 	d.checker = graph.NewConcurrentChecker(d, concurrencyLimits.Check, chunkSize)
 	d.expander = graph.NewConcurrentExpander(d)
 	d.lookupSubjectsHandler = graph.NewConcurrentLookupSubjects(d, concurrencyLimits.LookupSubjects, chunkSize)
-	d.lookupResourcesHandler2 = graph.NewCursoredLookupResources2(d, d, concurrencyLimits.LookupResources, chunkSize)
+	d.lookupResourcesHandler2 = graph.NewCursoredLookupResources2(d, d, typeSet, concurrencyLimits.LookupResources, chunkSize)
 
 	return d
 }
 
 // NewDispatcher creates a dispatcher that consults with the graph and redispatches subproblems to
 // the provided redispatcher.
-func NewDispatcher(redispatcher dispatch.Dispatcher, concurrencyLimits ConcurrencyLimits, dispatchChunkSize uint16) dispatch.Dispatcher {
+func NewDispatcher(redispatcher dispatch.Dispatcher, typeSet *caveattypes.TypeSet, concurrencyLimits ConcurrencyLimits, dispatchChunkSize uint16) dispatch.Dispatcher {
 	concurrencyLimits = limitsOrDefaults(concurrencyLimits, defaultConcurrencyLimit)
 	chunkSize := dispatchChunkSize
 	if chunkSize == 0 {
@@ -116,7 +117,7 @@ func NewDispatcher(redispatcher dispatch.Dispatcher, concurrencyLimits Concurren
 	checker := graph.NewConcurrentChecker(redispatcher, concurrencyLimits.Check, chunkSize)
 	expander := graph.NewConcurrentExpander(redispatcher)
 	lookupSubjectsHandler := graph.NewConcurrentLookupSubjects(redispatcher, concurrencyLimits.LookupSubjects, chunkSize)
-	lookupResourcesHandler2 := graph.NewCursoredLookupResources2(redispatcher, redispatcher, concurrencyLimits.LookupResources, chunkSize)
+	lookupResourcesHandler2 := graph.NewCursoredLookupResources2(redispatcher, redispatcher, typeSet, concurrencyLimits.LookupResources, chunkSize)
 
 	return &localDispatcher{
 		checker:                 checker,