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

repl/repl.go

@@ -909,12 +909,12 @@ func (r *REPL) evalBody(ctx context.Context, compiler *ast.Compiler, input ast.V
 
 	if r.explain != explainOff {
 		tracebuf = topdown.NewBufferTracer()
-		args = append(args, rego.Tracer(tracebuf))
+		args = append(args, rego.QueryTracer(tracebuf))
 	}
 
 	if r.profiler {
 		prof = profiler.New()
-		args = append(args, rego.Tracer(prof))
+		args = append(args, rego.QueryTracer(prof))
 	}
 
 	eval := rego.New(args...)
@@ -966,7 +966,7 @@ func (r *REPL) evalPartial(ctx context.Context, compiler *ast.Compiler, input as
 		rego.ParsedQuery(body),
 		rego.ParsedInput(input),
 		rego.Metrics(r.metrics),
-		rego.Tracer(buf),
+		rego.QueryTracer(buf),
 		rego.Instrument(r.instrument),
 		rego.ParsedUnknowns(r.unknowns),
 		rego.Runtime(r.runtime),

server/server.go

@@ -688,7 +688,7 @@ func (s *Server) execQuery(ctx context.Context, r *http.Request, txn storage.Tra
 		rego.ParsedInput(input),
 		rego.Metrics(m),
 		rego.Instrument(includeInstrumentation),
-		rego.Tracer(buf),
+		rego.QueryTracer(buf),
 		rego.Runtime(s.runtime),
 		rego.UnsafeBuiltins(unsafeBuiltinsMap),
 	)
@@ -1069,7 +1069,7 @@ func (s *Server) v1CompilePost(w http.ResponseWriter, r *http.Request) {
 		rego.ParsedQuery(request.Query),
 		rego.ParsedInput(request.Input),
 		rego.ParsedUnknowns(request.Unknowns),
-		rego.Tracer(buf),
+		rego.QueryTracer(buf),
 		rego.Instrument(includeInstrumentation),
 		rego.Metrics(m),
 		rego.Runtime(s.runtime),
@@ -1183,7 +1183,7 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 			rego.ParsedInput(input),
 			rego.Query(stringPathToDataRef(urlPath).String()),
 			rego.Metrics(m),
-			rego.Tracer(buf),
+			rego.QueryTracer(buf),
 			rego.Instrument(includeInstrumentation),
 			rego.Runtime(s.runtime),
 			rego.UnsafeBuiltins(unsafeBuiltinsMap),
@@ -1204,7 +1204,7 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 		rego.EvalTransaction(txn),
 		rego.EvalParsedInput(input),
 		rego.EvalMetrics(m),
-		rego.EvalTracer(buf),
+		rego.EvalQueryTracer(buf),
 	)
 
 	m.Timer(metrics.ServerHandler).Stop()
@@ -1399,7 +1399,7 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 		rego.EvalTransaction(txn),
 		rego.EvalParsedInput(input),
 		rego.EvalMetrics(m),
-		rego.EvalTracer(buf),
+		rego.EvalQueryTracer(buf),
 	)
 
 	m.Timer(metrics.ServerHandler).Stop()
@@ -2128,7 +2128,7 @@ func (s *Server) getCompiler() *ast.Compiler {
 	return s.manager.GetCompiler()
 }
 
-func (s *Server) makeRego(ctx context.Context, partial bool, txn storage.Transaction, input ast.Value, urlPath string, m metrics.Metrics, instrument bool, tracer topdown.Tracer, opts []func(*rego.Rego)) (*rego.Rego, error) {
+func (s *Server) makeRego(ctx context.Context, partial bool, txn storage.Transaction, input ast.Value, urlPath string, m metrics.Metrics, instrument bool, tracer topdown.QueryTracer, opts []func(*rego.Rego)) (*rego.Rego, error) {
 	queryPath := stringPathToDataRef(urlPath).String()
 
 	opts = append(
@@ -2137,7 +2137,7 @@ func (s *Server) makeRego(ctx context.Context, partial bool, txn storage.Transac
 		rego.Query(queryPath),
 		rego.ParsedInput(input),
 		rego.Metrics(m),
-		rego.Tracer(tracer),
+		rego.QueryTracer(tracer),
 		rego.Instrument(instrument),
 		rego.Runtime(s.runtime),
 		rego.UnsafeBuiltins(unsafeBuiltinsMap),

tester/runner.go

@@ -103,7 +103,7 @@ type BenchmarkOptions struct {
 type Runner struct {
 	compiler    *ast.Compiler
 	store       storage.Store
-	cover       topdown.Tracer
+	cover       topdown.QueryTracer
 	trace       bool
 	runtime     *ast.Term
 	failureLine bool
@@ -133,11 +133,27 @@ func (r *Runner) SetStore(store storage.Store) *Runner {
 }
 
 // SetCoverageTracer sets the tracer to use to compute coverage.
+// Deprecated: Use SetCoverageQueryTracer instead.
 func (r *Runner) SetCoverageTracer(tracer topdown.Tracer) *Runner {
-	r.cover = tracer
-	if r.cover != nil {
-		r.trace = false
+	if tracer == nil {
+		return r
 	}
+	if qt, ok := tracer.(topdown.QueryTracer); ok {
+		r.cover = qt
+	} else {
+		r.cover = topdown.WrapLegacyTracer(tracer)
+	}
+	r.trace = false
+	return r
+}
+
+// SetCoverageQueryTracer sets the tracer to use to compute coverage.
+func (r *Runner) SetCoverageQueryTracer(tracer topdown.QueryTracer) *Runner {
+	if tracer == nil {
+		return r
+	}
+	r.cover = tracer
+	r.trace = false
 	return r
 }
 
@@ -362,7 +378,7 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 
 	var bufferTracer *topdown.BufferTracer
 	var bufFailureLineTracer *topdown.BufferTracer
-	var tracer topdown.Tracer
+	var tracer topdown.QueryTracer
 
 	if r.cover != nil {
 		tracer = r.cover
@@ -379,7 +395,7 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 		rego.Transaction(txn),
 		rego.Compiler(r.compiler),
 		rego.Query(rule.Path().String()),
-		rego.Tracer(tracer),
+		rego.QueryTracer(tracer),
 		rego.Runtime(r.runtime),
 	)
 

tester/runner_test.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/cover"
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/tester"
 	"github.com/open-policy-agent/opa/topdown"
@@ -90,11 +91,23 @@ func TestRunner_EnableFailureLine(t *testing.T) {
 }
 
 func TestRun(t *testing.T) {
-	testRun(t, false)
+	testRun(t, testRunConfig{})
 }
 
 func TestRunBenchmark(t *testing.T) {
-	testRun(t, true)
+	testRun(t, testRunConfig{bench: true})
+}
+
+func TestRunWithCoverage(t *testing.T) {
+	cov := cover.New()
+	modules := testRun(t, testRunConfig{coverTracer: cov})
+	report := cov.Report(modules)
+	if len(report.Files) != len(modules) {
+		t.Errorf("Expected %d files in coverage report, got %d", len(modules), len(report.Files))
+	}
+	if report.Coverage == 0 {
+		t.Error("Expected test coverage")
+	}
 }
 
 type expectedTestResult struct {
@@ -102,9 +115,15 @@ type expectedTestResult struct {
 	wantFail bool
 }
 
+type testRunConfig struct {
+	bench       bool
+	filter      string
+	coverTracer topdown.QueryTracer
+}
+
 type expectedTestResults map[[2]string]expectedTestResult
 
-func testRun(t *testing.T, bench bool) {
+func testRun(t *testing.T, conf testRunConfig) map[string]*ast.Module {
 	files := map[string]string{
 		"/a.rego": `package foo
 			allow { true }
@@ -137,13 +156,16 @@ func testRun(t *testing.T, bench bool) {
 		{"data.bar", "test_duplicate"}:     {false, false},
 	}
 
+	var modules map[string]*ast.Module
 	test.WithTempFS(files, func(d string) {
-		rs := doTestRunWithTmpDir(t, d, bench, "")
-		validateTestResults(t, tests, rs, bench)
+		var rs []*tester.Result
+		rs, modules = doTestRunWithTmpDir(t, d, conf)
+		validateTestResults(t, tests, rs, conf)
 	})
+	return modules
 }
 
-func doTestRunWithTmpDir(t *testing.T, dir string, bench bool, regex string) []*tester.Result {
+func doTestRunWithTmpDir(t *testing.T, dir string, conf testRunConfig) ([]*tester.Result, map[string]*ast.Module) {
 	t.Helper()
 
 	ctx := context.Background()
@@ -157,10 +179,15 @@ func doTestRunWithTmpDir(t *testing.T, dir string, bench bool, regex string) []*
 	txn := storage.NewTransactionOrDie(ctx, store)
 	defer store.Abort(ctx, txn)
 
-	runner := tester.NewRunner().SetStore(store).SetModules(modules).Filter(regex).SetTimeout(60 * time.Second)
+	runner := tester.NewRunner().
+		SetStore(store).
+		SetModules(modules).
+		Filter(conf.filter).
+		SetTimeout(60 * time.Second).
+		SetCoverageQueryTracer(conf.coverTracer)
 
 	var ch chan *tester.Result
-	if bench {
+	if conf.bench {
 		ch, err = runner.RunBenchmarks(ctx, txn, tester.BenchmarkOptions{})
 	} else {
 		ch, err = runner.RunTests(ctx, txn)
@@ -173,10 +200,10 @@ func doTestRunWithTmpDir(t *testing.T, dir string, bench bool, regex string) []*
 		rs = append(rs, r)
 	}
 
-	return rs
+	return rs, modules
 }
 
-func validateTestResults(t *testing.T, tests expectedTestResults, rs []*tester.Result, bench bool) {
+func validateTestResults(t *testing.T, tests expectedTestResults, rs []*tester.Result, conf testRunConfig) {
 	t.Helper()
 	seen := map[[2]string]struct{}{}
 	for i := range rs {
@@ -189,9 +216,9 @@ func validateTestResults(t *testing.T, tests expectedTestResults, rs []*tester.R
 			t.Errorf("Expected %+v for %v but got: %v", exp, k, rs[i])
 		} else {
 			// Test passed
-			if bench && rs[i].BenchmarkResult == nil {
+			if conf.bench && rs[i].BenchmarkResult == nil {
 				t.Errorf("Expected BenchmarkResult for test %v, got nil", k)
-			} else if !bench && rs[i].BenchmarkResult != nil {
+			} else if !conf.bench && rs[i].BenchmarkResult != nil {
 				t.Errorf("Unexpected BenchmarkResult for test %v, expected nil", k)
 			}
 		}
@@ -320,8 +347,9 @@ func TestRunWithFilterRegex(t *testing.T) {
 
 		for _, tc := range cases {
 			t.Run(tc.note, func(t *testing.T) {
-				rs := doTestRunWithTmpDir(t, d, false, tc.regex)
-				validateTestResults(t, tc.tests, rs, false)
+				conf := testRunConfig{filter: tc.regex}
+				rs, _ := doTestRunWithTmpDir(t, d, conf)
+				validateTestResults(t, tc.tests, rs, conf)
 			})
 		}
 	})

topdown/query.go

@@ -104,7 +104,7 @@ func (q *Query) WithInput(input *ast.Term) *Query {
 func (q *Query) WithTracer(tracer Tracer) *Query {
 	qt, ok := tracer.(QueryTracer)
 	if !ok {
-		qt = wrapLegacyTracer(tracer)
+		qt = WrapLegacyTracer(tracer)
 	}
 	return q.WithQueryTracer(qt)
 }

topdown/trace.go

@@ -170,7 +170,9 @@ func (l *legacyTracer) TraceEvent(evt Event) {
 	l.t.Trace(&evt)
 }
 
-func wrapLegacyTracer(tracer Tracer) QueryTracer {
+// WrapLegacyTracer will create a new QueryTracer which wraps an
+// older Tracer instance.
+func WrapLegacyTracer(tracer Tracer) QueryTracer {
 	return &legacyTracer{t: tracer}
 }
 

watch/watch.go

@@ -256,7 +256,7 @@ func (h *Handle) deliver() {
 				rego.Store(h.watcher.store),
 				rego.Query(h.query),
 				rego.Metrics(m),
-				rego.Tracer(t),
+				rego.QueryTracer(t),
 				rego.Instrument(h.instrument),
 				rego.Runtime(h.runtime),
 			)