Fix composable schema interactions with expiration feature (#2780)

Author: tstirrat15

pkg/composableschemadsl/compiler/compiler.go

@@ -11,7 +11,6 @@ import (
 	"github.com/authzed/spicedb/pkg/composableschemadsl/input"
 	"github.com/authzed/spicedb/pkg/composableschemadsl/parser"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
-	"github.com/authzed/spicedb/pkg/genutil/slicez"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
 
@@ -55,7 +54,7 @@ func (cs CompiledSchema) SourcePositionToRunePosition(source input.Source, posit
 type config struct {
 	skipValidation   bool
 	objectTypePrefix *string
-	allowedFlags     []string
+	allowedFlags     *mapz.Set[string]
 	caveatTypeSet    *caveattypes.TypeSet
 
 	// In an import context, this is the folder containing
@@ -90,9 +89,7 @@ const expirationFlag = "expiration"
 
 func DisallowExpirationFlag() Option {
 	return func(cfg *config) {
-		cfg.allowedFlags = slicez.Filter(cfg.allowedFlags, func(s string) bool {
-			return s != expirationFlag
-		})
+		cfg.allowedFlags.Delete(expirationFlag)
 	}
 }
 
@@ -109,11 +106,11 @@ type ObjectPrefixOption func(*config)
 // Compile compilers the input schema into a set of namespace definition protos.
 func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*CompiledSchema, error) {
 	cfg := &config{
-		allowedFlags: make([]string, 0, 1),
+		allowedFlags: mapz.NewSet[string](),
 	}
 
 	// Enable `expiration` flag by default.
-	cfg.allowedFlags = append(cfg.allowedFlags, expirationFlag)
+	cfg.allowedFlags.Add(expirationFlag)
 
 	prefix(cfg) // required option
 
@@ -146,6 +143,7 @@ func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*Co
 		schemaString:       schema.SchemaString,
 		skipValidate:       cfg.skipValidation,
 		allowedFlags:       cfg.allowedFlags,
+		enabledFlags:       mapz.NewSet[string](),
 		existingNames:      mapz.NewSet[string](),
 		compiledPartials:   initialCompiledPartials,
 		unresolvedPartials: mapz.NewMultiMap[string, *dslNode](),

pkg/composableschemadsl/compiler/compiler_test.go

@@ -1254,17 +1254,23 @@ func TestCompile(t *testing.T) {
 			},
 		},
 		{
-			"duplicate use pragmas",
+			"wildcard relation with expiration trait",
 			withTenantPrefix,
-			`
-			use expiration
-			use expiration
-
+			`use expiration
+			
 			definition simple {
-				relation viewer: user with expiration
+				relation viewer: user:* with expiration
 			}`,
-			`found duplicate use flag`,
-			[]SchemaDefinition{},
+			"",
+			[]SchemaDefinition{
+				namespace.Namespace("sometenant/simple",
+					namespace.MustRelation("viewer", nil,
+						namespace.WithExpiration(
+							namespace.AllowedPublicNamespace("sometenant/user"),
+						),
+					),
+				),
+			},
 		},
 		{
 			"expiration use without use expiration",

pkg/composableschemadsl/compiler/importer-test/expiration-usage/dependency.zed

@@ -0,0 +1 @@
+use expiration

pkg/composableschemadsl/compiler/importer-test/expiration-usage/expected.zed

@@ -0,0 +1,9 @@
+use expiration
+
+definition test/user {}
+
+definition test/document {
+	relation viewer: test/user with expiration | test/user
+	relation public_reader: test/user:* with expiration | test/user:*
+	permission can_view = viewer + public_reader
+}

pkg/composableschemadsl/compiler/importer-test/expiration-usage/root.zed

@@ -0,0 +1,9 @@
+import "./dependency.zed"
+
+definition test/user {}
+
+definition test/document {
+  relation viewer: test/user with expiration | test/user
+  relation public_reader: test/user:* with expiration | test/user:*
+  permission can_view = viewer + public_reader
+}

pkg/composableschemadsl/compiler/importer-test/multiple-use-directives/expected.zed

@@ -0,0 +1,9 @@
+use expiration
+
+definition user {
+	relation timeout: user with expiration
+}
+
+definition resource {
+	relation tempuser: user with expiration
+}

pkg/composableschemadsl/compiler/importer-test/multiple-use-directives/root.zed

@@ -0,0 +1,7 @@
+use expiration
+
+import "./user.zed"
+
+definition resource {
+  relation tempuser: user with expiration
+}

pkg/composableschemadsl/compiler/importer-test/multiple-use-directives/user.zed

@@ -0,0 +1,5 @@
+use expiration
+
+definition user {
+  relation timeout: user with expiration
+}

pkg/composableschemadsl/compiler/importer_test.go

@@ -61,6 +61,8 @@ func TestImporter(t *testing.T) {
 		{"nested local import with transitive hop", "nested-local-with-hop"},
 		{"nested local two layers deep import", "nested-two-layer-local"},
 		{"diamond-shaped imports are fine", "diamond-shaped"},
+		{"multiple use directives are fine", "multiple-use-directives"},
+		{"expiration works correctly across multiple files", "expiration-usage"},
 	}
 
 	for _, test := range importerTests {

pkg/composableschemadsl/compiler/translator.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"path/filepath"
-	"slices"
 	"strings"
 
 	"github.com/ccoveille/go-safecast/v2"
@@ -27,8 +26,8 @@ type translationContext struct {
 	mapper           input.PositionMapper
 	schemaString     string
 	skipValidate     bool
-	allowedFlags     []string
-	enabledFlags     []string
+	allowedFlags     *mapz.Set[string]
+	enabledFlags     *mapz.Set[string]
 	existingNames    *mapz.Set[string]
 	caveatTypeSet    *caveattypes.TypeSet
 
@@ -668,29 +667,6 @@ func translateSpecificTypeReference(tctx *translationContext, typeRefNode *dslNo
 		return nil, typeRefNode.Errorf("%w", err)
 	}
 
-	if typeRefNode.Has(dslshape.NodeSpecificReferencePredicateWildcard) {
-		ref := &core.AllowedRelation{
-			Namespace: nspath,
-			RelationOrWildcard: &core.AllowedRelation_PublicWildcard_{
-				PublicWildcard: &core.AllowedRelation_PublicWildcard{},
-			},
-		}
-
-		err = addWithCaveats(tctx, typeRefNode, ref)
-		if err != nil {
-			return nil, typeRefNode.Errorf("invalid caveat: %w", err)
-		}
-
-		if !tctx.skipValidate {
-			if err := ref.Validate(); err != nil {
-				return nil, typeRefNode.Errorf("invalid type relation: %w", err)
-			}
-		}
-
-		ref.SourcePosition = getSourcePosition(typeRefNode, tctx.mapper)
-		return ref, nil
-	}
-
 	relationName := Ellipsis
 	if typeRefNode.Has(dslshape.NodeSpecificReferencePredicateRelation) {
 		relationName, err = typeRefNode.GetString(dslshape.NodeSpecificReferencePredicateRelation)
@@ -706,42 +682,57 @@ func translateSpecificTypeReference(tctx *translationContext, typeRefNode *dslNo
 		},
 	}
 
+	if typeRefNode.Has(dslshape.NodeSpecificReferencePredicateWildcard) {
+		ref.RelationOrWildcard = &core.AllowedRelation_PublicWildcard_{
+			PublicWildcard: &core.AllowedRelation_PublicWildcard{},
+		}
+	}
+
 	// Add the caveat(s), if any.
 	err = addWithCaveats(tctx, typeRefNode, ref)
 	if err != nil {
 		return nil, typeRefNode.Errorf("invalid caveat: %w", err)
 	}
 
 	// Add the expiration trait, if any.
+	err = addWithExpiration(tctx, typeRefNode, ref)
+	if err != nil {
+		return nil, typeRefNode.Errorf("invalid expiration: %w", err)
+	}
+
+	if !tctx.skipValidate {
+		if err := ref.Validate(); err != nil {
+			return nil, typeRefNode.Errorf("invalid type relation: %w", err)
+		}
+	}
+
+	ref.SourcePosition = getSourcePosition(typeRefNode, tctx.mapper)
+	return ref, nil
+}
+
+func addWithExpiration(tctx *translationContext, typeRefNode *dslNode, ref *core.AllowedRelation) error {
 	if traitNode, err := typeRefNode.Lookup(dslshape.NodeSpecificReferencePredicateTrait); err == nil {
 		traitName, err := traitNode.GetString(dslshape.NodeTraitPredicateTrait)
 		if err != nil {
-			return nil, typeRefNode.Errorf("invalid trait: %w", err)
+			return err
 		}
 
 		if traitName != "expiration" {
-			return nil, typeRefNode.Errorf("invalid trait: %s", traitName)
+			return fmt.Errorf("invalid trait: %s", traitName)
 		}
 
-		if !slices.Contains(tctx.allowedFlags, "expiration") {
-			return nil, typeRefNode.Errorf("expiration trait is not allowed")
+		if !tctx.allowedFlags.Has("expiration") {
+			return errors.New("expiration trait is not allowed")
 		}
 
-		if !slices.Contains(tctx.enabledFlags, "expiration") {
-			return nil, typeRefNode.Errorf("expiration flag is not enabled; add `use expiration` to top of file")
+		if !tctx.enabledFlags.Has("expiration") {
+			return errors.New("expiration flag is not enabled; add `use expiration` to top of file")
 		}
 
 		ref.RequiredExpiration = &core.ExpirationTrait{}
 	}
 
-	if !tctx.skipValidate {
-		if err := ref.Validate(); err != nil {
-			return nil, typeRefNode.Errorf("invalid type relation: %w", err)
-		}
-	}
-
-	ref.SourcePosition = getSourcePosition(typeRefNode, tctx.mapper)
-	return ref, nil
+	return nil
 }
 
 func addWithCaveats(tctx *translationContext, typeRefNode *dslNode, ref *core.AllowedRelation) error {
@@ -946,9 +937,9 @@ func translateUseFlag(tctx *translationContext, useFlagNode *dslNode) error {
 	if err != nil {
 		return err
 	}
-	if slices.Contains(tctx.enabledFlags, flagName) {
-		return useFlagNode.Errorf("found duplicate use flag: %s", flagName)
-	}
-	tctx.enabledFlags = append(tctx.enabledFlags, flagName)
+	// NOTE: we're okay with multiple instances of a given `use` directive in
+	// composable schemas, because each file may declare it separately
+	// and that should be valid.
+	tctx.enabledFlags.Add(flagName)
 	return nil
 }