refactor creation of prepared queries

Before this refactoring only the PreparedQuery options could be passed when using this plugin as a library. There are two more options (prepare options and evaluation options) that can be passed when using the OPA library directly.

This change breaks the interface but also allows for extensibility of creating the prepared queries in the future.

Signed-off-by: Magnus Jungsbluth <[email protected]>

envoyauth/evaluation.go

@@ -3,13 +3,11 @@ package envoyauth
 import (
 	"context"
 	"fmt"
-	"sync"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/config"
 	"github.com/open-policy-agent/opa/logging"
-	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/topdown/builtins"
@@ -25,17 +23,21 @@ type EvalContext interface {
 	Store() storage.Store
 	Compiler() *ast.Compiler
 	Runtime() *ast.Term
-	PreparedQueryDoOnce() *sync.Once
 	InterQueryBuiltinCache() iCache.InterQueryCache
-	PreparedQuery() *rego.PreparedEvalQuery
-	SetPreparedQuery(*rego.PreparedEvalQuery)
 	Logger() logging.Logger
 	Config() *config.Config
 	DistributedTracing() tracing.Options
+	CreatePreparedQueryOnce(opts PrepareQueryOpts) (*rego.PreparedEvalQuery, error)
+}
+
+// PrepareQueryOpts - Options to prepare a Rego query to be passed to the CreatePreparedQueryOnce method
+type PrepareQueryOpts struct {
+	Opts        []func(*rego.Rego)
+	PrepareOpts []rego.PrepareOption
 }
 
 // Eval - Evaluates an input against a provided EvalContext and yields result
-func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result *EvalResult, opts ...func(*rego.Rego)) error {
+func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result *EvalResult, evalOpts ...rego.EvalOption) error {
 	var err error
 	logger := evalContext.Logger()
 
@@ -64,7 +66,19 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result
 		"txn":   result.TxnID,
 	}).Debug("Executing policy query.")
 
-	err = constructPreparedQuery(evalContext, result.Txn, result.Metrics, opts)
+	pq, err := evalContext.CreatePreparedQueryOnce(
+		PrepareQueryOpts{
+			Opts: []func(*rego.Rego){
+				rego.Metrics(result.Metrics),
+				rego.ParsedQuery(evalContext.ParsedQuery()),
+				rego.Compiler(evalContext.Compiler()),
+				rego.Store(evalContext.Store()),
+				rego.Transaction(result.Txn),
+				rego.Runtime(evalContext.Runtime()),
+				rego.EnablePrintStatements(true),
+				rego.DistributedTracingOpts(evalContext.DistributedTracing()),
+			},
+		})
 	if err != nil {
 		return err
 	}
@@ -76,9 +90,7 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result
 		ndbCache = builtins.NDBCache{}
 	}
 
-	var rs rego.ResultSet
-	rs, err = evalContext.PreparedQuery().Eval(
-		ctx,
+	evalOpts = append(evalOpts,
 		rego.EvalParsedInput(input),
 		rego.EvalTransaction(result.Txn),
 		rego.EvalMetrics(result.Metrics),
@@ -87,6 +99,12 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result
 		rego.EvalNDBuiltinCache(ndbCache),
 	)
 
+	var rs rego.ResultSet
+	rs, err = pq.Eval(
+		ctx,
+		evalOpts...,
+	)
+
 	switch {
 	case err != nil:
 		return err
@@ -101,28 +119,6 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result
 	return nil
 }
 
-func constructPreparedQuery(evalContext EvalContext, txn storage.Transaction, m metrics.Metrics, opts []func(*rego.Rego)) error {
-	var err error
-	var pq rego.PreparedEvalQuery
-	evalContext.PreparedQueryDoOnce().Do(func() {
-		opts = append(opts,
-			rego.Metrics(m),
-			rego.ParsedQuery(evalContext.ParsedQuery()),
-			rego.Compiler(evalContext.Compiler()),
-			rego.Store(evalContext.Store()),
-			rego.Transaction(txn),
-			rego.Runtime(evalContext.Runtime()),
-			rego.EnablePrintStatements(true),
-			rego.DistributedTracingOpts(evalContext.DistributedTracing()),
-		)
-
-		pq, err = rego.New(opts...).PrepareForEval(context.Background())
-		evalContext.SetPreparedQuery(&pq)
-	})
-
-	return err
-}
-
 func getRevision(ctx context.Context, store storage.Store, txn storage.Transaction, result *EvalResult) error {
 	revisions := map[string]string{}
 

envoyauth/evaluation_test.go

@@ -254,6 +254,7 @@ type mockExtAuthzGrpcServer struct {
 	manager                *plugins.Manager
 	preparedQuery          *rego.PreparedEvalQuery
 	preparedQueryDoOnce    *sync.Once
+	preparedQueryErr       error
 	distributedTracingOpts tracing.Options
 }
 
@@ -301,6 +302,17 @@ func (m *mockExtAuthzGrpcServer) DistributedTracing() tracing.Options {
 	return m.distributedTracingOpts
 }
 
+func (m *mockExtAuthzGrpcServer) CreatePreparedQueryOnce(opts PrepareQueryOpts) (*rego.PreparedEvalQuery, error) {
+	m.preparedQueryDoOnce.Do(func() {
+		pq, err := rego.New(opts.Opts...).PrepareForEval(context.Background())
+
+		m.preparedQuery = &pq
+		m.preparedQueryErr = err
+	})
+
+	return m.preparedQuery, m.preparedQueryErr
+}
+
 type testPlugin struct {
 	events []logs.EventV1
 }

internal/internal.go

@@ -225,6 +225,7 @@ type envoyExtAuthzGrpcServer struct {
 	manager                *plugins.Manager
 	preparedQuery          *rego.PreparedEvalQuery
 	preparedQueryDoOnce    *sync.Once
+	preparedQueryErr       error
 	interQueryBuiltinCache iCache.InterQueryCache
 	distributedTracingOpts tracing.Options
 	metricAuthzDuration    prometheus.HistogramVec
@@ -279,6 +280,17 @@ func (p *envoyExtAuthzGrpcServer) DistributedTracing() tracing.Options {
 	return p.distributedTracingOpts
 }
 
+func (p *envoyExtAuthzGrpcServer) CreatePreparedQueryOnce(opts envoyauth.PrepareQueryOpts) (*rego.PreparedEvalQuery, error) {
+	p.preparedQueryDoOnce.Do(func() {
+		pq, err := rego.New(opts.Opts...).PrepareForEval(context.Background())
+
+		p.preparedQuery = &pq
+		p.preparedQueryErr = err
+	})
+
+	return p.preparedQuery, p.preparedQueryErr
+}
+
 func (p *envoyExtAuthzGrpcServer) Start(ctx context.Context) error {
 	p.manager.UpdatePluginStatus(PluginName, &plugins.Status{State: plugins.StateNotReady})
 	go p.listen()