rego: Plumb the newer QueryTracer through Rego API's

This will deprecate the older API's that used the `topdown.Tracer` in
favor of the newer `topdown.QueryTracer` interface. Usages of the old
API have been swapped, although some of the testing is left with them
to ensure we still support them (until we remove the deprecated API).

Signed-off-by: Patrick East <east.patrick@gmail.com>

CHANGELOG.md

@@ -18,6 +18,13 @@ failure via the status API and error logging. For more information see https://o
 Thanks to @ashish246 who co-designed the feature and provided valuable input to the development process with his
 proof-of-concept [#1757](https://github.com/open-policy-agent/opa/issues/1757).
 
+### Backwards Compatibility
+
+* The `rego.Tracer` and `rego.EvalTracer` API's have been deprecated in favor of
+  the newer `rego.QueryTracer` and `rego.EvalQueryTracer` API.
+* The `tester.Runner#SetCoverageTracer` API has been deprecated in favor of the
+  newer `test.Runner#SetCoverageQueryTracer` API.
+
 ## 0.21.1
 
 This release fixes [#2497](https://github.com/open-policy-agent/opa/issues/2497) where the comprehension indexing optimization produced incorrect results for nested comprehensions that close over variables in the outer scope. This issue only affects policies containing nested comprehensions that are recognized by the indexer (which is a relatively small percentage).

cmd/eval.go

@@ -410,7 +410,7 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) {
 
 	if params.explain != nil && params.explain.String() != explainModeOff {
 		tracer = topdown.NewBufferTracer()
-		evalArgs = append(evalArgs, rego.EvalTracer(tracer))
+		evalArgs = append(evalArgs, rego.EvalQueryTracer(tracer))
 	}
 
 	if params.disableIndexing {
@@ -434,7 +434,7 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) {
 	var p *profiler.Profiler
 	if params.profile {
 		p = profiler.New()
-		evalArgs = append(evalArgs, rego.EvalTracer(p))
+		evalArgs = append(evalArgs, rego.EvalQueryTracer(p))
 	}
 
 	if params.partial {
@@ -447,7 +447,7 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) {
 
 	if params.coverage {
 		c = cover.New()
-		evalArgs = append(evalArgs, rego.EvalTracer(c))
+		evalArgs = append(evalArgs, rego.EvalQueryTracer(c))
 	}
 
 	eval := rego.New(regoArgs...)

cmd/test.go

@@ -189,7 +189,7 @@ func opaTest(args []string) int {
 	}
 
 	var cov *cover.Cover
-	var coverTracer topdown.Tracer
+	var coverTracer topdown.QueryTracer
 
 	if testParams.coverage {
 		if testParams.benchmark {
@@ -204,7 +204,7 @@ func opaTest(args []string) int {
 		SetCompiler(compiler).
 		SetStore(store).
 		EnableTracing(testParams.verbose).
-		SetCoverageTracer(coverTracer).
+		SetCoverageQueryTracer(coverTracer).
 		EnableFailureLine(testParams.failureLine).
 		SetRuntime(info).
 		SetModules(modules).

cmd/test_test.go

@@ -180,7 +180,7 @@ func failTrace(t *testing.T) []*topdown.Event {
 	_, err := rego.New(
 		rego.Module("test.rego", mod),
 		rego.Trace(true),
-		rego.Tracer(tracer),
+		rego.QueryTracer(tracer),
 		rego.Query("data.testing.test_p"),
 	).Eval(context.Background())
 

cover/cover_bench_test.go

@@ -45,7 +45,7 @@ func BenchmarkCoverBigLocalVar(b *testing.B) {
 
 				for i := 0; i < b.N; i++ {
 					b.StartTimer()
-					_, err = pq.Eval(ctx, rego.EvalTracer(cover))
+					_, err = pq.Eval(ctx, rego.EvalQueryTracer(cover))
 					b.StopTimer()
 
 					if err != nil {

cover/cover_test.go

@@ -57,7 +57,7 @@ p {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(cover),
+		rego.QueryTracer(cover),
 	)
 
 	ctx := context.Background()

profiler/profiler_bench_test.go

@@ -45,7 +45,7 @@ func BenchmarkProfilerBigLocalVar(b *testing.B) {
 
 				for i := 0; i < b.N; i++ {
 					b.StartTimer()
-					_, err = pq.Eval(ctx, rego.EvalTracer(profiler))
+					_, err = pq.Eval(ctx, rego.EvalQueryTracer(profiler))
 					b.StopTimer()
 
 					if err != nil {

profiler/profiler_test.go

@@ -53,7 +53,7 @@ baz {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -125,7 +125,7 @@ func TestProfileCheckExprDuration(t *testing.T) {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -192,7 +192,7 @@ baz {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -246,7 +246,7 @@ baz {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -307,7 +307,7 @@ baz {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -374,7 +374,7 @@ baz {
 	eval := rego.New(
 		rego.Module("test.rego", module),
 		rego.Query("data.test.foo"),
-		rego.Tracer(profiler),
+		rego.QueryTracer(profiler),
 	)
 
 	ctx := context.Background()
@@ -444,7 +444,7 @@ allowed_operations = [
 		t.Fatal(err)
 	}
 
-	_, err = pq.Eval(ctx, rego.EvalTracer(profiler))
+	_, err = pq.Eval(ctx, rego.EvalQueryTracer(profiler))
 	if err != nil {
 		t.Fatal(err)
 	}

rego/example_test.go

@@ -593,7 +593,7 @@ func ExampleRego_Eval_tracer() {
 	// Create very simple query that binds a single variable and provides a tracer.
 	rego := rego.New(
 		rego.Query("x = 1"),
-		rego.Tracer(buf),
+		rego.QueryTracer(buf),
 	)
 
 	// Run evaluation.

rego/rego.go

@@ -88,7 +88,7 @@ type EvalContext struct {
 	instrument       bool
 	instrumentation  *topdown.Instrumentation
 	partialNamespace string
-	tracers          []topdown.Tracer
+	queryTracers     []topdown.QueryTracer
 	compiledQuery    compiledQuery
 	unknowns         []string
 	disableInlining  []ast.Ref
@@ -137,10 +137,20 @@ func EvalInstrument(instrument bool) EvalOption {
 }
 
 // EvalTracer configures a tracer for a Prepared Query's evaluation
+// Deprecated: Use EvalQueryTracer instead.
 func EvalTracer(tracer topdown.Tracer) EvalOption {
 	return func(e *EvalContext) {
 		if tracer != nil {
-			e.tracers = append(e.tracers, tracer)
+			e.queryTracers = append(e.queryTracers, topdown.WrapLegacyTracer(tracer))
+		}
+	}
+}
+
+// EvalQueryTracer configures a tracer for a Prepared Query's evaluation
+func EvalQueryTracer(tracer topdown.QueryTracer) EvalOption {
+	return func(e *EvalContext) {
+		if tracer != nil {
+			e.queryTracers = append(e.queryTracers, tracer)
 		}
 	}
 }
@@ -216,7 +226,7 @@ func (pq preparedQuery) newEvalContext(ctx context.Context, options []EvalOption
 		instrument:       false,
 		instrumentation:  nil,
 		partialNamespace: pq.r.partialNamespace,
-		tracers:          nil,
+		queryTracers:     nil,
 		unknowns:         pq.r.unknowns,
 		parsedUnknowns:   pq.r.parsedUnknowns,
 		compiledQuery:    compiledQuery{},
@@ -457,7 +467,7 @@ type Rego struct {
 	ownStore               bool
 	txn                    storage.Transaction
 	metrics                metrics.Metrics
-	tracers                []topdown.Tracer
+	queryTracers           []topdown.QueryTracer
 	tracebuf               *topdown.BufferTracer
 	trace                  bool
 	instrumentation        *topdown.Instrumentation
@@ -885,10 +895,20 @@ func Trace(yes bool) func(r *Rego) {
 }
 
 // Tracer returns an argument that adds a query tracer to r.
+// Deprecated: Use QueryTracer instead.
 func Tracer(t topdown.Tracer) func(r *Rego) {
 	return func(r *Rego) {
 		if t != nil {
-			r.tracers = append(r.tracers, t)
+			r.queryTracers = append(r.queryTracers, topdown.WrapLegacyTracer(t))
+		}
+	}
+}
+
+// QueryTracer returns an argument that adds a query tracer to r.
+func QueryTracer(t topdown.QueryTracer) func(r *Rego) {
+	return func(r *Rego) {
+		if t != nil {
+			r.queryTracers = append(r.queryTracers, t)
 		}
 	}
 }
@@ -976,7 +996,7 @@ func New(options ...func(r *Rego)) *Rego {
 
 	if r.trace {
 		r.tracebuf = topdown.NewBufferTracer()
-		r.tracers = append(r.tracers, r.tracebuf)
+		r.queryTracers = append(r.queryTracers, r.tracebuf)
 	}
 
 	if r.partialNamespace == "" {
@@ -1007,8 +1027,8 @@ func (r *Rego) Eval(ctx context.Context) (ResultSet, error) {
 		EvalInstrument(r.instrument),
 	}
 
-	for _, t := range r.tracers {
-		evalArgs = append(evalArgs, EvalTracer(t))
+	for _, qt := range r.queryTracers {
+		evalArgs = append(evalArgs, EvalQueryTracer(qt))
 	}
 
 	rs, err := pq.Eval(ctx, evalArgs...)
@@ -1074,8 +1094,8 @@ func (r *Rego) Partial(ctx context.Context) (*PartialQueries, error) {
 		EvalInstrument(r.instrument),
 	}
 
-	for _, t := range r.tracers {
-		evalArgs = append(evalArgs, EvalTracer(t))
+	for _, t := range r.queryTracers {
+		evalArgs = append(evalArgs, EvalQueryTracer(t))
 	}
 
 	pqs, err := pq.Partial(ctx, evalArgs...)
@@ -1632,8 +1652,8 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
 		WithRuntime(r.runtime).
 		WithIndexing(ectx.indexing)
 
-	for i := range ectx.tracers {
-		q = q.WithTracer(ectx.tracers[i])
+	for i := range ectx.queryTracers {
+		q = q.WithQueryTracer(ectx.queryTracers[i])
 	}
 
 	if ectx.parsedInput != nil {
@@ -1716,7 +1736,7 @@ func (r *Rego) partialResult(ctx context.Context, pCfg *PrepareConfig) (PartialR
 		metrics:          r.metrics,
 		txn:              r.txn,
 		partialNamespace: r.partialNamespace,
-		tracers:          r.tracers,
+		queryTracers:     r.queryTracers,
 		compiledQuery:    r.compiledQueries[partialResultQueryType],
 		instrumentation:  r.instrumentation,
 		indexing:         true,
@@ -1820,8 +1840,8 @@ func (r *Rego) partial(ctx context.Context, ectx *EvalContext) (*PartialQueries,
 		WithSkipPartialNamespace(r.skipPartialNamespace).
 		WithShallowInlining(r.shallowInlining)
 
-	for i := range ectx.tracers {
-		q = q.WithTracer(ectx.tracers[i])
+	for i := range ectx.queryTracers {
+		q = q.WithQueryTracer(ectx.queryTracers[i])
 	}
 
 	if ectx.parsedInput != nil {

rego/rego_test.go

@@ -465,6 +465,34 @@ func TestPreparedRegoTracerNoPropagate(t *testing.T) {
 	}
 }
 
+func TestPreparedRegoQueryTracerNoPropagate(t *testing.T) {
+	tracer := topdown.NewBufferTracer()
+	mod := `
+	package test
+
+	p = {
+		input.x == 10
+	}
+	`
+	pq, err := New(
+		Query("data"),
+		Module("foo.rego", mod),
+		QueryTracer(tracer),
+		Input(map[string]interface{}{"x": 10})).PrepareForEval(context.Background())
+	if err != nil {
+		t.Fatalf("unexpected error %s", err)
+	}
+
+	_, err = pq.Eval(context.Background()) // no EvalQueryTracer option
+	if err != nil {
+		t.Fatalf("unexpected error %s", err)
+	}
+
+	if len(*tracer) > 0 {
+		t.Fatal("expected 0 traces to be collected")
+	}
+}
+
 func TestRegoDisableIndexing(t *testing.T) {
 	tracer := topdown.NewBufferTracer()
 	mod := `
@@ -488,7 +516,7 @@ func TestRegoDisableIndexing(t *testing.T) {
 
 	_, err = pq.Eval(
 		context.Background(),
-		EvalTracer(tracer),
+		EvalQueryTracer(tracer),
 		EvalRuleIndexing(false),
 		EvalInput(map[string]interface{}{"x": 10}),
 	)
@@ -1073,6 +1101,46 @@ func TestPreparedPartialResultWithTracer(t *testing.T) {
 	}
 }
 
+func TestPreparedPartialResultWithQueryTracer(t *testing.T) {
+	mod := `
+	package test
+	default p = false
+	p {
+		input.x = 1
+	}
+	`
+	r := New(
+		Query("data.test.p == true"),
+		Module("test.rego", mod),
+	)
+
+	tracer := topdown.NewBufferTracer()
+
+	ctx := context.Background()
+	pq, err := r.PrepareForPartial(ctx)
+	if err != nil {
+		t.Fatalf("unexpected error from Rego.PrepareForPartial(): %s", err.Error())
+	}
+
+	pqs, err := pq.Partial(ctx, EvalQueryTracer(tracer))
+	if err != nil {
+		t.Fatalf("unexpected error from PreparedEvalQuery.Partial(): %s", err.Error())
+	}
+
+	expectedQuery := "input.x = 1"
+	if len(pqs.Queries) != 1 {
+		t.Errorf("expected 1 query but found %d: %+v", len(pqs.Queries), pqs)
+	}
+	if pqs.Queries[0].String() != expectedQuery {
+		t.Errorf("unexpected query in result, expected='%s' found='%s'",
+			expectedQuery, pqs.Queries[0].String())
+	}
+
+	if len(*tracer) == 0 {
+		t.Errorf("Expected buffer tracer to contain > 0 traces")
+	}
+}
+
 func TestPartialResultSetsValidConflictChecker(t *testing.T) {
 	mod := `
 	package test