From 6e73023974e7c7b1efc867c3b0f62788b5aceca1 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Tue, 26 Nov 2024 12:32:17 -0500
Subject: [PATCH] Add explicit option to enable expiration in schema

This will allow us to test without allowing the keyword otherwise
---
 pkg/schemadsl/compiler/compiler.go        | 6 ++++++
 pkg/schemadsl/compiler/compiler_test.go   | 2 +-
 pkg/schemadsl/compiler/translator.go      | 6 ++++++
 pkg/schemadsl/generator/generator_test.go | 2 +-
 4 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/pkg/schemadsl/compiler/compiler.go b/pkg/schemadsl/compiler/compiler.go
index 3511be9fb2..efa7f47715 100644
--- a/pkg/schemadsl/compiler/compiler.go
+++ b/pkg/schemadsl/compiler/compiler.go
@@ -52,6 +52,7 @@ func (cs CompiledSchema) SourcePositionToRunePosition(source input.Source, posit
 type config struct {
 	skipValidation   bool
 	objectTypePrefix *string
+	allowedFlags     []string
 }
 
 func SkipValidation() Option { return func(cfg *config) { cfg.skipValidation = true } }
@@ -68,6 +69,10 @@ func AllowUnprefixedObjectType() ObjectPrefixOption {
 	return func(cfg *config) { cfg.objectTypePrefix = new(string) }
 }
 
+func AllowExpirationUseFlag() Option {
+	return func(cfg *config) { cfg.allowedFlags = append(cfg.allowedFlags, "expiration") }
+}
+
 type Option func(*config)
 
 type ObjectPrefixOption func(*config)
@@ -94,6 +99,7 @@ func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*Co
 		mapper:           mapper,
 		schemaString:     schema.SchemaString,
 		skipValidate:     cfg.skipValidation,
+		allowedFlags:     cfg.allowedFlags,
 	}, root)
 	if err != nil {
 		var errorWithNode errorWithNode
diff --git a/pkg/schemadsl/compiler/compiler_test.go b/pkg/schemadsl/compiler/compiler_test.go
index 114b8ffef4..78fa8a579d 100644
--- a/pkg/schemadsl/compiler/compiler_test.go
+++ b/pkg/schemadsl/compiler/compiler_test.go
@@ -1032,7 +1032,7 @@ func TestCompile(t *testing.T) {
 			require := require.New(t)
 			compiled, err := Compile(InputSchema{
 				input.Source(test.name), test.input,
-			}, test.objectPrefix)
+			}, test.objectPrefix, AllowExpirationUseFlag())
 
 			if test.expectedError != "" {
 				require.Error(err)
diff --git a/pkg/schemadsl/compiler/translator.go b/pkg/schemadsl/compiler/translator.go
index 7e69fcc0f8..c1421ba67f 100644
--- a/pkg/schemadsl/compiler/translator.go
+++ b/pkg/schemadsl/compiler/translator.go
@@ -3,6 +3,7 @@ package compiler
 import (
 	"bufio"
 	"fmt"
+	"slices"
 	"strings"
 
 	"github.com/ccoveille/go-safecast"
@@ -22,6 +23,7 @@ type translationContext struct {
 	mapper           input.PositionMapper
 	schemaString     string
 	skipValidate     bool
+	allowedFlags     []string
 }
 
 func (tctx translationContext) prefixedPath(definitionName string) (string, error) {
@@ -650,6 +652,10 @@ func translateSpecificTypeReference(tctx translationContext, typeRefNode *dslNod
 			return nil, typeRefNode.Errorf("invalid trait: %s", traitName)
 		}
 
+		if !slices.Contains(tctx.allowedFlags, "expiration") {
+			return nil, typeRefNode.Errorf("expiration trait is not allowed")
+		}
+
 		ref.RequiredExpiration = &core.ExpirationTrait{}
 	}
 
diff --git a/pkg/schemadsl/generator/generator_test.go b/pkg/schemadsl/generator/generator_test.go
index a9cbc89495..82bf4fb78e 100644
--- a/pkg/schemadsl/generator/generator_test.go
+++ b/pkg/schemadsl/generator/generator_test.go
@@ -405,7 +405,7 @@ definition document {
 			compiled, err := compiler.Compile(compiler.InputSchema{
 				Source:       input.Source(test.name),
 				SchemaString: test.input,
-			}, compiler.AllowUnprefixedObjectType())
+			}, compiler.AllowUnprefixedObjectType(), compiler.AllowExpirationUseFlag())
 			require.NoError(err)
 
 			source, _, err := GenerateSchema(compiled.OrderedDefinitions)