rego: Allow per-eval GenerateJSON function (#8690)

The Rego API's `GenerateJSON` function allows clients to provide custom
logic for transforming an original AST result into whatever format they
may need. Previously this could only be set on the Rego object directly,
meaning that a single prepared query would have to use the same function
for all evaluations. This change adds the option to additionally set an
`EvalGenerateJSON function scoped to individual evaluations, making it
easier to reuse a single prepared query even when the shape of the
result is determined dynamically, by input data, in-policy routing, etc.

Signed-off-by: Anders Eknert <anders.eknert@apple.com>

v1/rego/rego.go

@@ -100,6 +100,7 @@ type EvalContext struct {
 	parsedInput                 ast.Value
 	metrics                     metrics.Metrics
 	txn                         storage.Transaction
+	generateJSON                func(*ast.Term, *EvalContext) (any, error)
 	instrument                  bool
 	instrumentation             *topdown.Instrumentation
 	partialNamespace            string
@@ -227,6 +228,15 @@ func EvalTransaction(txn storage.Transaction) EvalOption {
 	}
 }
 
+// EvalGenerateJSON sets the AST to JSON converter for an evaluation. When set, this
+// option takes precedence over [GenerateJSON] set on the Rego object from e.g. a prepared
+// query, allowing individual evaluations to customize how the result is transformed.
+func EvalGenerateJSON(f func(*ast.Term, *EvalContext) (any, error)) EvalOption {
+	return func(e *EvalContext) {
+		e.generateJSON = f
+	}
+}
+
 // EvalInstrument enables or disables instrumenting for a Prepared Query's evaluation
 func EvalInstrument(instrument bool) EvalOption {
 	return func(e *EvalContext) {
@@ -1330,7 +1340,10 @@ func Target(t string) func(r *Rego) {
 	}
 }
 
-// GenerateJSON sets the AST to JSON converter for the results.
+// GenerateJSON sets the AST to JSON converter to use for results. This will have any evaluationn on a Rego
+// object use the provided function for conversion. Use [EvalGenerateJSON] if you want to set a converter
+// function for individual evaluations, which will take precedence over GenerateJSON set on the Rego object,
+// (i.e. by this function) for the scope of that evaluation.
 func GenerateJSON(f func(*ast.Term, *EvalContext) (any, error)) func(r *Rego) {
 	return func(r *Rego) {
 		r.generateJSON = f
@@ -2462,6 +2475,11 @@ func (r *Rego) valueToQueryResult(res ast.Value, ectx *EvalContext) (ResultSet,
 func (r *Rego) generateResult(qr topdown.QueryResult, ectx *EvalContext) (Result, error) {
 	rewritten := ectx.compiledQuery.compiler.RewrittenVars()
 
+	generateJSON := r.generateJSON
+	if ectx.generateJSON != nil {
+		generateJSON = ectx.generateJSON
+	}
+
 	result := newResult()
 	for k, term := range qr {
 		if rw, ok := rewritten[k]; ok {
@@ -2471,7 +2489,7 @@ func (r *Rego) generateResult(qr topdown.QueryResult, ectx *EvalContext) (Result
 			continue
 		}
 
-		v, err := r.generateJSON(term, ectx)
+		v, err := generateJSON(term, ectx)
 		if err != nil {
 			return result, err
 		}
@@ -2485,7 +2503,7 @@ func (r *Rego) generateResult(qr topdown.QueryResult, ectx *EvalContext) (Result
 		}
 
 		if k, ok := r.capture[expr]; ok {
-			v, err := r.generateJSON(qr[k], ectx)
+			v, err := generateJSON(qr[k], ectx)
 			if err != nil {
 				return result, err
 			}

v1/rego/rego_test.go

@@ -3173,6 +3173,35 @@ func TestGenerateJSON(t *testing.T) {
 	assertEval(t, r, `[["converted-input"]]`)
 }
 
+func TestEvalGenerateJSON(t *testing.T) {
+	// Given a prepared query with a GenerateJSON function provided
+	r := New(Query("input"), Input("original-input"), GenerateJSON(generateJSONResult("prepare-generate")))
+	pq, err := r.PrepareForEval(t.Context())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// EvalGenerateJSON takes precedence over the GenerateJSON set on the prepared query
+	rs, err := pq.Eval(t.Context(), EvalGenerateJSON(generateJSONResult("eval-generate")))
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertResultSet(t, rs, `[["eval-generate"]]`)
+
+	// When EvalGenerateJSON is not provided, the GenerateJSON from the prepared query is used
+	rs, err = pq.Eval(t.Context())
+	if err != nil {
+		t.Fatal(err)
+	}
+	assertResultSet(t, rs, `[["prepare-generate"]]`)
+}
+
+func generateJSONResult(s string) func(*ast.Term, *EvalContext) (any, error) {
+	return func(*ast.Term, *EvalContext) (any, error) {
+		return s, nil
+	}
+}
+
 func TestRegoLazyObjDefault(t *testing.T) {
 	foo := map[string]any{"foo": "bar", "other": 1}
 	store := inmem.NewFromObjectWithOpts(map[string]any{