feat(engine): add query metadata to evaluation results (#1061)

* feat(engine): add query metadata to evaluation results Add query information to result metadata to enable output formatters to utilize the originating query in their output. This helps track which policy rule triggered each result. The change: - Adds query metadata to both string and map return types - Ensures metadata map is properly initialized - Adds comprehensive tests covering single and multiple results Signed-off-by: Ville Vesilehto <[email protected]> * test: split checks in TestQueryMetadata Check msg and meta in separate blocks, with separate err Signed-off-by: Ville Vesilehto <[email protected]> --------- Signed-off-by: Ville Vesilehto <[email protected]>
open-policy-agent · Feb 11, 2025 · eacba23 · eacba23
1 parent 6da5673
commit eacba23
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 0 deletions.
diff --git a/policy/engine.go b/policy/engine.go
@@ -500,6 +500,9 @@ func (e *Engine) query(ctx context.Context, input interface{}, query string) (ou
 				case string:
 					result := output.Result{
 						Message: val,
+						Metadata: map[string]any{
+							"query": query,
+						},
 					}
 					results = append(results, result)
 
@@ -510,6 +513,9 @@ func (e *Engine) query(ctx context.Context, input interface{}, query string) (ou
 						return output.QueryResult{}, fmt.Errorf("new result: %w", err)
 					}
 
+					// Safe to set as Metadata map is initialized by NewResult
+					result.Metadata["query"] = query
+
 					results = append(results, result)
 				}
 			}

diff --git a/policy/engine_test.go b/policy/engine_test.go
@@ -573,3 +573,133 @@ deny[msg] { msg := "denied" }`),
 		})
 	}
 }
+
+func TestQueryMetadata(t *testing.T) {
+	type want struct {
+		msg  string
+		meta map[string]any
+	}
+
+	tests := []struct {
+		name   string
+		policy []byte
+		query  string
+		want   []want
+	}{
+		{
+			name: "string return type",
+			policy: []byte(`package main
+deny[msg] {
+	msg := "simple denial"
+}`),
+			query: "data.main.deny",
+			want: []want{
+				{
+					msg:  "simple denial",
+					meta: map[string]any{"query": "data.main.deny"},
+				},
+			},
+		},
+		{
+			name: "map return type",
+			policy: []byte(`package main
+violation[result] {
+	result := {
+		"msg": "violation with metadata",
+		"severity": "high"
+	}
+}`),
+			query: "data.main.violation",
+			want: []want{
+				{
+					msg: "violation with metadata",
+					meta: map[string]any{
+						"query":    "data.main.violation",
+						"severity": "high",
+					},
+				},
+			},
+		},
+		{
+			name: "multiple results",
+			policy: []byte(`package main
+deny[msg] {
+	msg := "first denial"
+}
+deny[msg] {
+	msg := "second denial"
+}
+violation[result] {
+	result := {
+		"msg": "violation one",
+		"severity": "high"
+	}
+}
+violation[result] {
+	result := {
+		"msg": "violation two",
+		"severity": "low"
+	}
+}`),
+			query: "data.main.deny",
+			want: []want{
+				{
+					msg:  "first denial",
+					meta: map[string]any{"query": "data.main.deny"},
+				},
+				{
+					msg:  "second denial",
+					meta: map[string]any{"query": "data.main.deny"},
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+
+			files := fstest.MapFS{
+				"policy.rego": &fstest.MapFile{
+					Data: tt.policy,
+				},
+			}
+
+			l := loader.NewFileLoader().WithFS(files)
+			pols, err := l.All([]string{"policy.rego"})
+			if err != nil {
+				t.Fatalf("Load policies: %v", err)
+			}
+
+			engine := Engine{
+				modules:  pols.ParsedModules(),
+				compiler: ast.NewCompiler().WithEnablePrintStatements(true),
+			}
+			engine.compiler.Compile(engine.modules)
+			if engine.compiler.Failed() {
+				t.Fatalf("Compiler error: %v", engine.compiler.Errors)
+			}
+
+			result, err := engine.query(ctx, nil, tt.query)
+			if err != nil {
+				t.Fatalf("Query error: %v", err)
+			}
+
+			if len(result.Results) != len(tt.want) {
+				t.Fatalf("got %d results, want %d", len(result.Results), len(tt.want))
+			}
+
+			for i, got := range result.Results {
+				want := tt.want[i]
+				if got.Message != want.msg {
+					t.Errorf("result[%d]: got Message=%q, want Message=%q",
+						i, got.Message, want.msg)
+				}
+				if !reflect.DeepEqual(got.Metadata, want.meta) {
+					t.Errorf("result[%d]: got Metadata=%v, want Metadata=%v",
+						i, got.Metadata, want.meta)
+				}
+			}
+		})
+	}
+}