Add support to the schema generator for the new expiration trait

pkg/schemadsl/generator/generator.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/authzed/spicedb/pkg/caveats"
 	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	"github.com/authzed/spicedb/pkg/graph"
 	"github.com/authzed/spicedb/pkg/namespace"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -26,6 +27,8 @@ const MaxSingleLineCommentLength = 70 // 80 - the comment parts and some padding
 // GenerateSchema generates a DSL view of the given schema.
 func GenerateSchema(definitions []compiler.SchemaDefinition) (string, bool, error) {
 	generated := make([]string, 0, len(definitions))
+	flags := mapz.NewSet[string]()
+
 	result := true
 	for _, definition := range definitions {
 		switch def := definition.(type) {
@@ -39,19 +42,29 @@ func GenerateSchema(definitions []compiler.SchemaDefinition) (string, bool, erro
 			generated = append(generated, generatedCaveat)
 
 		case *core.NamespaceDefinition:
-			generatedSchema, ok, err := GenerateSource(def)
+			generatedSchema, defFlags, ok, err := generateDefinitionSource(def)
 			if err != nil {
 				return "", false, err
 			}
 
 			result = result && ok
 			generated = append(generated, generatedSchema)
+			flags.Extend(defFlags)
 
 		default:
 			return "", false, spiceerrors.MustBugf("unknown type of definition %T in GenerateSchema", def)
 		}
 	}
 
+	if !flags.IsEmpty() {
+		flagsSlice := flags.AsSlice()
+		sort.Strings(flagsSlice)
+
+		for _, flag := range flagsSlice {
+			generated = append([]string{"use " + flag}, generated...)
+		}
+	}
+
 	return strings.Join(generated, "\n\n"), result, nil
 }
 
@@ -74,19 +87,25 @@ func GenerateCaveatSource(caveat *core.CaveatDefinition) (string, bool, error) {
 
 // GenerateSource generates a DSL view of the given namespace definition.
 func GenerateSource(namespace *core.NamespaceDefinition) (string, bool, error) {
+	source, _, ok, err := generateDefinitionSource(namespace)
+	return source, ok, err
+}
+
+func generateDefinitionSource(namespace *core.NamespaceDefinition) (string, []string, bool, error) {
 	generator := &sourceGenerator{
 		indentationLevel: 0,
 		hasNewline:       true,
 		hasBlankline:     true,
 		hasNewScope:      true,
+		flags:            mapz.NewSet[string](),
 	}
 
 	err := generator.emitNamespace(namespace)
 	if err != nil {
-		return "", false, err
+		return "", nil, false, err
 	}
 
-	return generator.buf.String(), !generator.hasIssue, nil
+	return generator.buf.String(), generator.flags.AsSlice(), !generator.hasIssue, nil
 }
 
 // GenerateRelationSource generates a DSL view of the given relation definition.
@@ -237,9 +256,25 @@ func (sg *sourceGenerator) emitAllowedRelation(allowedRelation *core.AllowedRela
 	if allowedRelation.GetPublicWildcard() != nil {
 		sg.append(":*")
 	}
-	if allowedRelation.GetRequiredCaveat() != nil {
+
+	hasExpirationTrait := allowedRelation.GetRequiredExpiration() != nil
+	hasCaveat := allowedRelation.GetRequiredCaveat() != nil
+
+	if hasExpirationTrait || hasCaveat {
 		sg.append(" with ")
-		sg.append(allowedRelation.RequiredCaveat.CaveatName)
+		if hasCaveat {
+			sg.append(allowedRelation.RequiredCaveat.CaveatName)
+		}
+
+		if hasExpirationTrait {
+			sg.flags.Add("expiration")
+
+			if hasCaveat {
+				sg.append(" and ")
+			}
+
+			sg.append("expiration")
+		}
 	}
 }
 

pkg/schemadsl/generator/generator_impl.go

@@ -2,16 +2,19 @@ package generator
 
 import (
 	"strings"
+
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
 )
 
 type sourceGenerator struct {
-	buf                strings.Builder // The buffer for the new source code.
-	indentationLevel   int             // The current indentation level.
-	hasNewline         bool            // Whether there is a newline at the end of the buffer.
-	hasBlankline       bool            // Whether there is a blank line at the end of the buffer.
-	hasIssue           bool            // Whether there is a translation issue.
-	hasNewScope        bool            // Whether there is a new scope at the end of the buffer.
-	existingLineLength int             // Length of the existing line.
+	buf                strings.Builder   // The buffer for the new source code.
+	indentationLevel   int               // The current indentation level.
+	hasNewline         bool              // Whether there is a newline at the end of the buffer.
+	hasBlankline       bool              // Whether there is a blank line at the end of the buffer.
+	hasIssue           bool              // Whether there is a translation issue.
+	hasNewScope        bool              // Whether there is a new scope at the end of the buffer.
+	existingLineLength int               // Length of the existing line.
+	flags              *mapz.Set[string] // The flags added while generating.
 }
 
 // ensureBlankLineOrNewScope ensures that there is a blank line or new scope at the tail of the buffer. If not,

pkg/schemadsl/generator/generator_test.go

@@ -304,7 +304,6 @@ definition foos/test {
 	relation somerel: foos/bars
 }`,
 		},
-
 		{
 			"full example",
 			`
@@ -360,6 +359,41 @@ definition foos/document {
 }`,
 			`definition document {
 	permission first = rela->relb + relc.any(reld) + rele.all(relf)
+}`,
+		},
+		{
+			"expiration caveat",
+			`definition document{
+				relation viewer: user with expiration
+		}`,
+			`definition document {
+	relation viewer: user with expiration
+}`,
+		},
+		{
+			"expiration trait",
+			`use expiration
+			
+			definition document{
+				relation viewer: user with expiration
+				relation editor: user with somecaveat and expiration
+		}`,
+			`use expiration
+
+definition document {
+	relation viewer: user with expiration
+	relation editor: user with somecaveat and expiration
+}`,
+		},
+		{
+			"unused expiration flag",
+			`use expiration
+			
+			definition document{
+				relation viewer: user
+		}`,
+			`definition document {
+	relation viewer: user
 }`,
 		},
 	}