From 91027d1e2659985f12cafcce7dc7288fe1b01171 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Fri, 14 Jan 2022 12:31:27 -0800 Subject: [PATCH] cmd+compile: add new plan target This commit adds a new "plan" target to the compile package and build command. This targets lets users compile out query plans into bundles. This way the compiled query plan can be consumed by transpilers/interpreters outside of OPA. Fixes #4133 Signed-off-by: Torin Sandall --- cmd/build.go | 2 +- compile/compile.go | 105 ++++++++++++++++++++++++++++++---------- compile/compile_test.go | 22 +++++++++ 3 files changed, 102 insertions(+), 27 deletions(-) diff --git a/cmd/build.go b/cmd/build.go index ea6a591d85..4e1fe5ff4e 100644 --- a/cmd/build.go +++ b/cmd/build.go @@ -46,7 +46,7 @@ type buildParams struct { func newBuildParams() buildParams { return buildParams{ capabilities: newcapabilitiesFlag(), - target: util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}), + target: util.NewEnumFlag(compile.TargetRego, compile.Targets), } } diff --git a/compile/compile.go b/compile/compile.go index 9d7fdca94d..bce2286da9 100644 --- a/compile/compile.go +++ b/compile/compile.go @@ -8,6 +8,7 @@ package compile import ( "bytes" "context" + "encoding/json" "fmt" "io" "regexp" @@ -40,15 +41,21 @@ const ( // TargetWasm is an alternative target that compiles the policy into a wasm // module instead of Rego. The target supports base documents. TargetWasm = "wasm" + + // TargetPlan is an altertive target that compiles the policy into an + // imperative query plan that can be further transpiled or interpreted. + TargetPlan = "plan" ) -const wasmResultVar = ast.Var("result") - -var validTargets = map[string]struct{}{ - TargetRego: {}, - TargetWasm: {}, +// Targets contains the list of targets supported by the compiler. +var Targets = []string{ + TargetRego, + TargetWasm, + TargetPlan, } +const resultVar = ast.Var("result") + // Compiler implements bundle compilation and linking. type Compiler struct { capabilities *ast.Capabilities // the capabilities that compiled policies may require @@ -63,6 +70,7 @@ type Compiler struct { 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 @@ -187,10 +195,28 @@ func (c *Compiler) Build(ctx context.Context) error { return err } - if c.target == TargetWasm { + switch c.target { + case TargetWasm: if err := c.compileWasm(ctx); err != nil { return err } + case TargetPlan: + if err := c.compilePlan(ctx); err != nil { + return err + } + + bs, err := json.Marshal(c.policy) + if err != nil { + return err + } + + c.bundle.PlanModules = append(c.bundle.PlanModules, bundle.PlanModuleFile{ + Path: bundle.PlanFile, + URL: bundle.PlanFile, + Raw: bs, + }) + case TargetRego: + // nop } if c.revision != nil { @@ -220,7 +246,15 @@ func (c *Compiler) init() error { c.capabilities = ast.CapabilitiesForThisVersion() } - if _, ok := validTargets[c.target]; !ok { + var found bool + for _, t := range Targets { + if c.target == t { + found = true + break + } + } + + if !found { return fmt.Errorf("invalid target %q", c.target) } @@ -238,8 +272,15 @@ func (c *Compiler) init() error { return errors.New("bundle optimizations require at least one entrypoint") } - if c.target == TargetWasm && len(c.entrypointrefs) == 0 { - return errors.New("wasm compilation requires at least one entrypoint") + switch c.target { + case TargetWasm: + if len(c.entrypointrefs) == 0 { + return errors.New("wasm compilation requires at least one entrypoint") + } + case TargetPlan: + if len(c.entrypointrefs) == 0 { + return errors.New("plan compilation requires at least one entrypoint") + } } return nil @@ -342,7 +383,7 @@ func (c *Compiler) optimize(ctx context.Context) error { return nil } -func (c *Compiler) compileWasm(ctx context.Context) error { +func (c *Compiler) compilePlan(ctx context.Context) error { // Lazily compile the modules if needed. If optimizations were run, the // AST compiler will not be set because the default target does not require it. @@ -384,7 +425,7 @@ func (c *Compiler) compileWasm(ctx context.Context) error { } // Create query sets for each of the entrypoints. - resultSym := ast.NewTerm(wasmResultVar) + resultSym := ast.NewTerm(resultVar) queries := make([]planner.QuerySet, len(c.entrypointrefs)) for i := range c.entrypointrefs { @@ -414,7 +455,32 @@ func (c *Compiler) compileWasm(ctx context.Context) error { builtins[bi.Name] = bi } + // Plan the query sets. + p := planner.New(). + WithQueries(queries). + WithModules(modules). + WithBuiltinDecls(builtins). + WithDebug(c.debug.Writer()) + policy, err := p.Plan() + if err != nil { + return err + } + + // dump policy IR (if "debug" wasn't requested, debug.Witer will discard it) + err = ir.Pretty(c.debug.Writer(), policy) + if err != nil { + return err + } + + c.policy = policy + + return nil +} + +func (c *Compiler) compileWasm(ctx context.Context) error { + compiler := wasm.New() + found := false have := compiler.ABIVersion() if c.capabilities.WasmABIVersions == nil { // discern nil from len=0 @@ -434,25 +500,12 @@ func (c *Compiler) compileWasm(ctx context.Context) error { ) } - // Plan the query sets. - p := planner.New(). - WithQueries(queries). - WithModules(modules). - WithBuiltinDecls(builtins). - WithDebug(c.debug.Writer()) - policy, err := p.Plan() - if err != nil { - return err - } - - // dump policy IR (if "debug" wasn't requested, debug.Witer will discard it) - err = ir.Pretty(c.debug.Writer(), policy) - if err != nil { + if err := c.compilePlan(ctx); err != nil { return err } // Compile the policy into a wasm binary. - m, err := compiler.WithPolicy(policy).WithDebug(c.debug.Writer()).Compile() + m, err := compiler.WithPolicy(c.policy).WithDebug(c.debug.Writer()).Compile() if err != nil { return err } diff --git a/compile/compile_test.go b/compile/compile_test.go index 4b6524dfbe..800516c9f0 100644 --- a/compile/compile_test.go +++ b/compile/compile_test.go @@ -660,6 +660,28 @@ func ensureEntrypointRemoved(t *testing.T, b *bundle.Bundle, e string) { } } +func TestCompilerPlanTarget(t *testing.T) { + files := map[string]string{ + "test.rego": `package test + + p = 7 + q = p+1`, + } + + test.WithTempFS(files, func(root string) { + + compiler := New().WithPaths(root).WithTarget("plan").WithEntrypoints("test/p", "test/q") + err := compiler.Build(context.Background()) + if err != nil { + t.Fatal(err) + } + + if len(compiler.bundle.PlanModules) == 0 { + t.Fatal("expected to find compiled plan module") + } + }) +} + func TestCompilerSetRevision(t *testing.T) { files := map[string]string{ "test.rego": `package test