add compiler.Compiler and compiler.optimizer WithEnablePrintStatements option and pass through to ast.Compiler

When using the compiler.Compiler; rego "print" statements get removed and there is no way on the API to enable them. So the PrintHook is broken. Signed-off-by: Kevin St. Pierre <kevin@styra.com>
2026-08-12 19:32:48 -06:00 · 2022-07-22 09:05:33 -04:00
parent 9f9fbb91e9
commit 646c8418ae
2 changed files with 107 additions and 33 deletions
@@ -57,24 +57,25 @@ const resultVar = ast.Var("result")

 // Compiler implements bundle compilation and linking.
 type Compiler struct {
-	capabilities      *ast.Capabilities          // the capabilities that compiled policies may require
-	bundle            *bundle.Bundle             // the bundle that the compiler operates on
-	revision          *string                    // the revision to set on the output bundle
-	asBundle          bool                       // whether to assume bundle layout on file loading or not
-	filter            loader.Filter              // filter to apply to file loader
-	paths             []string                   // file paths to load. TODO(tsandall): add support for supplying readers for embedded users.
-	entrypoints       orderedStringSet           // policy entrypoints required for optimization and certain targets
-	optimizationLevel int                        // how aggressive should optimization be
-	target            string                     // target type (wasm, rego, etc.)
-	output            *io.Writer                 // output stream to write bundle to
-	entrypointrefs    []*ast.Term                // validated entrypoints computed from default decision or manually supplied entrypoints
-	compiler          *ast.Compiler              // rego ast compiler used for semantic checks and rewriting
-	policy            *ir.Policy                 // planner output when wasm or plan targets are enabled
-	debug             debug.Debug                // optionally outputs debug information produced during build
-	bvc               *bundle.VerificationConfig // represents the key configuration used to verify a signed bundle
-	bsc               *bundle.SigningConfig      // represents the key configuration used to generate a signed bundle
-	keyID             string                     // represents the name of the default key used to verify a signed bundle
-	metadata          *map[string]interface{}    // represents additional data included in .manifest file
+	capabilities          *ast.Capabilities          // the capabilities that compiled policies may require
+	bundle                *bundle.Bundle             // the bundle that the compiler operates on
+	revision              *string                    // the revision to set on the output bundle
+	asBundle              bool                       // whether to assume bundle layout on file loading or not
+	filter                loader.Filter              // filter to apply to file loader
+	paths                 []string                   // file paths to load. TODO(tsandall): add support for supplying readers for embedded users.
+	entrypoints           orderedStringSet           // policy entrypoints required for optimization and certain targets
+	optimizationLevel     int                        // how aggressive should optimization be
+	target                string                     // target type (wasm, rego, etc.)
+	output                *io.Writer                 // output stream to write bundle to
+	entrypointrefs        []*ast.Term                // validated entrypoints computed from default decision or manually supplied entrypoints
+	compiler              *ast.Compiler              // rego ast compiler used for semantic checks and rewriting
+	policy                *ir.Policy                 // planner output when wasm or plan targets are enabled
+	debug                 debug.Debug                // optionally outputs debug information produced during build
+	enablePrintStatements bool                       // optionally enable rego print statements
+	bvc                   *bundle.VerificationConfig // represents the key configuration used to verify a signed bundle
+	bsc                   *bundle.SigningConfig      // represents the key configuration used to generate a signed bundle
+	keyID                 string                     // represents the name of the default key used to verify a signed bundle
+	metadata              *map[string]interface{}    // represents additional data included in .manifest file
 }

 // New returns a new compiler instance that can be invoked.
@@ -135,6 +136,14 @@ func (c *Compiler) WithDebug(sink io.Writer) *Compiler {
 	return c
 }

+// WithEnablePrintStatements enables print statements inside of modules compiled
+// by the compiler. If print statements are not enabled, calls to print() are
+// erased at compile-time.
+func (c *Compiler) WithEnablePrintStatements(yes bool) *Compiler {
+	c.enablePrintStatements = yes
+	return c
+}
+
 // WithPaths adds input filepaths to read policy and data from.
 func (c *Compiler) WithPaths(p ...string) *Compiler {
 	c.paths = append(c.paths, p...)
@@ -374,14 +383,15 @@ func (c *Compiler) optimize(ctx context.Context) error {

 	if c.optimizationLevel <= 0 {
 		var err error
-		c.compiler, err = compile(c.capabilities, c.bundle, c.debug)
+		c.compiler, err = compile(c.capabilities, c.bundle, c.debug, c.enablePrintStatements)
 		return err
 	}

 	o := newOptimizer(c.capabilities, c.bundle).
 		WithEntrypoints(c.entrypointrefs).
 		WithDebug(c.debug.Writer()).
-		WithShallowInlining(c.optimizationLevel <= 1)
+		WithShallowInlining(c.optimizationLevel <= 1).
+		WithEnablePrintStatements(c.enablePrintStatements)

 	err := o.Do(ctx)
 	if err != nil {
@@ -399,7 +409,7 @@ func (c *Compiler) compilePlan(ctx context.Context) error {
 	// AST compiler will not be set because the default target does not require it.
 	if c.compiler == nil {
 		var err error
-		c.compiler, err = compile(c.capabilities, c.bundle, c.debug)
+		c.compiler, err = compile(c.capabilities, c.bundle, c.debug, c.enablePrintStatements)
 		if err != nil {
 			return err
 		}
@@ -615,15 +625,16 @@ func (err undefinedEntrypointErr) Error() string {
 }

 type optimizer struct {
-	capabilities    *ast.Capabilities
-	bundle          *bundle.Bundle
-	compiler        *ast.Compiler
-	entrypoints     []*ast.Term
-	nsprefix        string
-	resultsymprefix string
-	outputprefix    string
-	shallow         bool
-	debug           debug.Debug
+	capabilities          *ast.Capabilities
+	bundle                *bundle.Bundle
+	compiler              *ast.Compiler
+	entrypoints           []*ast.Term
+	nsprefix              string
+	resultsymprefix       string
+	outputprefix          string
+	shallow               bool
+	debug                 debug.Debug
+	enablePrintStatements bool
 }

 func newOptimizer(c *ast.Capabilities, b *bundle.Bundle) *optimizer {
@@ -644,6 +655,11 @@ func (o *optimizer) WithDebug(sink io.Writer) *optimizer {
 	return o
 }

+func (o *optimizer) WithEnablePrintStatements(yes bool) *optimizer {
+	o.enablePrintStatements = yes
+	return o
+}
+
 func (o *optimizer) WithEntrypoints(es []*ast.Term) *optimizer {
 	o.entrypoints = es
 	return o
@@ -683,7 +699,7 @@ func (o *optimizer) Do(ctx context.Context) error {
 	for i, e := range o.entrypoints {

 		var err error
-		o.compiler, err = compile(o.capabilities, o.bundle, o.debug)
+		o.compiler, err = compile(o.capabilities, o.bundle, o.debug, o.enablePrintStatements)
 		if err != nil {
 			return err
 		}
@@ -939,7 +955,7 @@ func (o *optimizer) getSupportModuleFilename(used map[string]int, module *ast.Mo

 var safePathPattern = regexp.MustCompile(`^[\w-_/]+$`)

-func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug) (*ast.Compiler, error) {
+func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug, enablePrintStatements bool) (*ast.Compiler, error) {

 	modules := map[string]*ast.Module{}

@@ -951,7 +967,7 @@ func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug) (*ast.Compi
 		modules[mf.URL] = mf.Parsed
 	}

-	compiler := ast.NewCompiler().WithCapabilities(c).WithDebug(dbg.Writer())
+	compiler := ast.NewCompiler().WithCapabilities(c).WithDebug(dbg.Writer()).WithEnablePrintStatements(enablePrintStatements)
 	compiler.Compile(modules)

 	if compiler.Failed() {
@@ -215,6 +215,64 @@ func TestCompilerLoadFilesystem(t *testing.T) {
 	})
 }

+func TestCompilerLoadFilesystemWithEnablePrintStatementsFalse(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+			package test
+
+                        allow { print(1) }
+		`,
+		"data.json": `
+			{"b1": {"k": "v"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+		compiler := New().
+			WithPaths(root).
+			WithTarget("plan").WithEntrypoints("test/allow").
+			WithEnablePrintStatements(false)
+
+		if err := compiler.Build(context.Background()); err != nil {
+			t.Fatal(err)
+		}
+
+		bundle := compiler.Bundle()
+
+		if strings.Contains(string(bundle.PlanModules[0].Raw), "internal.print") {
+			t.Fatalf("output different than expected:\n\ngot: %v\n\nfound: internal.print", string(bundle.PlanModules[0].Raw))
+		}
+	})
+}
+
+func TestCompilerLoadFilesystemWithEnablePrintStatementsTrue(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+			package test
+
+                        allow { print(1) }
+		`,
+		"data.json": `
+			{"b1": {"k": "v"}}`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+		compiler := New().
+			WithPaths(root).
+			WithTarget("plan").WithEntrypoints("test/allow").
+			WithEnablePrintStatements(true)
+
+		if err := compiler.Build(context.Background()); err != nil {
+			t.Fatal(err)
+		}
+
+		bundle := compiler.Bundle()
+
+		if !strings.Contains(string(bundle.PlanModules[0].Raw), "internal.print") {
+			t.Fatalf("output different than expected:\n\ngot: %v\n\nmissing: internal.print", string(bundle.PlanModules[0].Raw))
+		}
+	})
+}
+
 func TestCompilerLoadHonorsFilter(t *testing.T) {
 	files := map[string]string{
 		"test.rego": `