server: wire in response/request metadata for compile handler

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

v1/server/compile_handler.go

@@ -88,6 +88,49 @@ type CompileFiltersRequestV1 struct {
 		TargetDialects  []string       `json:"targetDialects,omitempty"`
 		MaskRule        string         `json:"maskRule,omitempty"`
 	} `json:"options"`
+
+	// Metadata holds any additional top-level fields not defined in this struct.
+	// These fields are preserved during JSON unmarshaling, allowing wrapping
+	// projects to pass through custom key/value pairs (e.g. snapshot_id).
+	Metadata map[string]any `json:"-"`
+}
+
+var compileFiltersKnownKeys = map[string]bool{
+	"input": true, "query": true, "unknowns": true, "options": true,
+}
+
+func (r *CompileFiltersRequestV1) UnmarshalJSON(data []byte) error {
+	type Alias CompileFiltersRequestV1
+	aux := &struct {
+		*Alias
+	}{
+		Alias: (*Alias)(r),
+	}
+
+	if err := util.UnmarshalJSON(data, aux); err != nil {
+		return err
+	}
+
+	var raw map[string]json.RawMessage
+	if err := util.UnmarshalJSON(data, &raw); err != nil {
+		return err
+	}
+
+	for key, val := range raw {
+		if compileFiltersKnownKeys[key] {
+			continue
+		}
+		if r.Metadata == nil {
+			r.Metadata = make(map[string]any)
+		}
+		var v any
+		if err := util.UnmarshalJSON(val, &v); err != nil {
+			return err
+		}
+		r.Metadata[key] = v
+	}
+
+	return nil
 }
 
 type compileFiltersRequest struct {
@@ -107,6 +150,37 @@ type CompileResponseV1 struct {
 	Explanation types.TraceV1     `json:"explanation,omitempty"`
 	Metrics     types.MetricsV1   `json:"metrics,omitempty"`
 	Hints       []failtracer.Hint `json:"hints,omitempty"`
+
+	Metadata map[string]any `json:"-"`
+}
+
+var compileResponseReservedFields = map[string]bool{
+	"result": true, "explanation": true, "metrics": true, "hints": true,
+}
+
+func (r CompileResponseV1) MarshalJSON() ([]byte, error) {
+	type Alias CompileResponseV1
+	data, err := json.Marshal(Alias(r))
+	if err != nil {
+		return nil, err
+	}
+
+	if len(r.Metadata) == 0 {
+		return data, nil
+	}
+
+	var base map[string]any
+	if err := json.Unmarshal(data, &base); err != nil {
+		return nil, err
+	}
+
+	for key, val := range r.Metadata {
+		if !compileResponseReservedFields[key] {
+			base[key] = val
+		}
+	}
+
+	return json.Marshal(base)
 }
 
 func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
@@ -266,7 +340,8 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
 
 	qt := failtracer.New()
 
-	filters, err := preparedCompile.Compile(ctx,
+	respMetadata := map[string]any{}
+	evalOpts := []rego.EvalOption{
 		rego.EvalTransaction(txn),
 		rego.EvalParsedInput(request.Input),
 		rego.EvalPrintHook(s.manager.PrintHook()),
@@ -274,7 +349,13 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
 		rego.EvalInterQueryBuiltinCache(s.interQueryBuiltinCache),
 		rego.EvalInterQueryBuiltinValueCache(s.interQueryBuiltinValueCache),
 		rego.EvalQueryTracer(qt),
-	)
+		rego.EvalResponseMetadata(respMetadata),
+	}
+	if orig.Metadata != nil {
+		evalOpts = append(evalOpts, rego.EvalRequestMetadata(orig.Metadata))
+	}
+
+	filters, err := preparedCompile.Compile(ctx, evalOpts...)
 	if err != nil {
 		switch err := err.(type) {
 		case ast.Errors:
@@ -338,6 +419,11 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
 	}
 
 	m.Timer(metrics.ServerHandler).Stop()
+
+	if len(respMetadata) > 0 {
+		result.Metadata = respMetadata
+	}
+
 	fin(w, result, contentType, m, includeMetrics(r), includeInstrumentation, pretty(r))
 
 	unk := make([]string, len(unknowns))
@@ -353,6 +439,12 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
 		"type":      decisionLogType,
 		"mask_rule": maskingRule.String(),
 	}
+	if len(orig.Metadata) > 0 {
+		custom["request_metadata"] = orig.Metadata
+	}
+	if len(respMetadata) > 0 {
+		custom["response_metadata"] = respMetadata
+	}
 
 	if err := logger.Log(ctx, txn, urlPath, orig.Query, orig.Input, request.Input, result.Result, ndbCache, nil, m, nil, custom); err != nil {
 		writer.ErrorAuto(w, err)

v1/server/compile_handler_test.go

@@ -377,6 +377,263 @@ func TestCompileHandlerMaskingRules(t *testing.T) {
 	})
 }
 
+func TestCompileFiltersRequestUnmarshalMetadata(t *testing.T) {
+	t.Parallel()
+
+	body := `{
+		"input": {"user": "alice"},
+		"unknowns": ["input.resources"],
+		"options": {"maskRule": "masks"},
+		"snapshot_id": "t|2026-05-13T00:00:00Z|2026-05-13T00:10:00Z|s",
+		"com.example.opa/metadata": {"trace_id": "xyz-789"}
+	}`
+
+	var req CompileFiltersRequestV1
+	if err := json.Unmarshal([]byte(body), &req); err != nil {
+		t.Fatal(err)
+	}
+
+	if req.Input == nil {
+		t.Fatal("expected input")
+	}
+	if req.Unknowns == nil || len(*req.Unknowns) != 1 || (*req.Unknowns)[0] != "input.resources" {
+		t.Fatalf("unexpected unknowns: %v", req.Unknowns)
+	}
+	if req.Options.MaskRule != "masks" {
+		t.Fatalf("unexpected maskRule: %v", req.Options.MaskRule)
+	}
+
+	if req.Metadata == nil {
+		t.Fatal("expected metadata")
+	}
+	if req.Metadata["snapshot_id"] != "t|2026-05-13T00:00:00Z|2026-05-13T00:10:00Z|s" {
+		t.Fatalf("unexpected snapshot_id: %v", req.Metadata["snapshot_id"])
+	}
+	md, ok := req.Metadata["com.example.opa/metadata"].(map[string]any)
+	if !ok {
+		t.Fatal("expected com.example.opa/metadata in metadata")
+	}
+	if md["trace_id"] != "xyz-789" {
+		t.Fatalf("unexpected trace_id: %v", md["trace_id"])
+	}
+
+	// Known fields must not appear in metadata
+	for _, key := range []string{"input", "unknowns", "options", "query"} {
+		if _, ok := req.Metadata[key]; ok {
+			t.Fatalf("'%s' should not be in metadata", key)
+		}
+	}
+}
+
+func TestCompileResponseMarshalMetadata(t *testing.T) {
+	t.Parallel()
+
+	result := any("WHERE name = 'alice'")
+	resp := CompileResponseV1{
+		Result:   &result,
+		Metadata: map[string]any{"snapshot_id": "t|new-snapshot"},
+	}
+
+	data, err := json.Marshal(resp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var decoded map[string]any
+	if err := json.Unmarshal(data, &decoded); err != nil {
+		t.Fatal(err)
+	}
+
+	if decoded["snapshot_id"] != "t|new-snapshot" {
+		t.Fatalf("expected snapshot_id in response JSON, got %v", decoded["snapshot_id"])
+	}
+	if decoded["result"] == nil {
+		t.Fatal("expected result in response JSON")
+	}
+}
+
+func TestCompileResponseMarshalMetadataNoOverride(t *testing.T) {
+	t.Parallel()
+
+	result := any("WHERE name = 'alice'")
+	resp := CompileResponseV1{
+		Result: &result,
+		Metadata: map[string]any{
+			"result":      "should-not-override",
+			"snapshot_id": "t|valid",
+		},
+	}
+
+	data, err := json.Marshal(resp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var decoded map[string]any
+	if err := json.Unmarshal(data, &decoded); err != nil {
+		t.Fatal(err)
+	}
+
+	// "result" must not be overridden by metadata
+	if decoded["result"] != "WHERE name = 'alice'" {
+		t.Fatalf("result should not be overridden, got %v", decoded["result"])
+	}
+	if decoded["snapshot_id"] != "t|valid" {
+		t.Fatalf("expected snapshot_id, got %v", decoded["snapshot_id"])
+	}
+}
+
+func TestCompileResponseMarshalNoMetadata(t *testing.T) {
+	t.Parallel()
+
+	result := any("WHERE name = 'alice'")
+	resp := CompileResponseV1{
+		Result: &result,
+	}
+
+	data, err := json.Marshal(resp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var decoded map[string]any
+	if err := json.Unmarshal(data, &decoded); err != nil {
+		t.Fatal(err)
+	}
+
+	if _, ok := decoded["snapshot_id"]; ok {
+		t.Fatal("snapshot_id should not appear without metadata")
+	}
+}
+
+func TestCompileHandlerRequestMetadata(t *testing.T) {
+	t.Parallel()
+
+	rego := `package filters
+# METADATA
+# scope: document
+# compile:
+#   unknowns: [input.fruits]
+include if input.fruits.name == "apple"
+`
+
+	var logged *Info
+	f := setup(t, rego, nil)
+	f.server = f.server.WithDecisionLoggerWithErr(func(_ context.Context, info *Info) error {
+		logged = info
+		return nil
+	})
+
+	payload := map[string]any{
+		"input":                    map[string]any{"max": 1},
+		"snapshot_id":              "t|2026-05-13T00:00:00Z|2026-05-13T00:10:00Z|s",
+		"com.example.opa/metadata": map[string]any{"trace_id": "abc-123"},
+	}
+
+	req := evalReq(t, "filters/include", payload, "application/vnd.opa.sql.postgresql+json")
+	if err := f.executeRequest(req, http.StatusOK, ""); err != nil {
+		t.Fatal(err)
+	}
+
+	if logged == nil {
+		t.Fatal("expected decision log entry")
+	}
+	if logged.Custom == nil {
+		t.Fatal("expected Custom in decision log")
+	}
+
+	incoming, ok := logged.Custom["request_metadata"].(map[string]any)
+	if !ok {
+		t.Fatalf("expected request_metadata in Custom, got %v", logged.Custom)
+	}
+
+	if incoming["snapshot_id"] != "t|2026-05-13T00:00:00Z|2026-05-13T00:10:00Z|s" {
+		t.Fatalf("expected snapshot_id in request_metadata, got %v", incoming["snapshot_id"])
+	}
+
+	md, ok := incoming["com.example.opa/metadata"].(map[string]any)
+	if !ok {
+		t.Fatal("expected com.example.opa/metadata in request_metadata")
+	}
+	if md["trace_id"] != "abc-123" {
+		t.Fatalf("expected trace_id='abc-123', got %v", md["trace_id"])
+	}
+
+	// Known fields must not leak into metadata
+	for _, key := range []string{"input", "unknowns", "options", "query"} {
+		if _, ok := incoming[key]; ok {
+			t.Fatalf("'%s' should not appear in request_metadata", key)
+		}
+	}
+
+	if _, ok := logged.Custom["response_metadata"]; ok {
+		t.Fatal("response_metadata should be absent when nothing populates it")
+	}
+}
+
+func TestCompileHandlerResponseMetadata(t *testing.T) {
+	t.Parallel()
+
+	rego := `package filters
+# METADATA
+# scope: document
+# compile:
+#   unknowns: [input.fruits]
+include if {
+    test.set_outgoing()
+    input.fruits.name == "apple"
+}
+`
+
+	var logged *Info
+	f := setup(t, rego, nil)
+	f.server = f.server.WithDecisionLoggerWithErr(func(_ context.Context, info *Info) error {
+		logged = info
+		return nil
+	})
+
+	payload := map[string]any{
+		"input":       map[string]any{"max": 1},
+		"snapshot_id": "t|2026-05-13T00:00:00Z|2026-05-13T00:10:00Z|s",
+	}
+
+	req := evalReq(t, "filters/include", payload, "application/vnd.opa.sql.postgresql+json")
+	if err := f.executeRequest(req, http.StatusOK, ""); err != nil {
+		t.Fatal(err)
+	}
+
+	// Check the HTTP response body contains response metadata
+	var respBody map[string]any
+	if err := json.NewDecoder(f.recorder.Result().Body).Decode(&respBody); err != nil {
+		t.Fatal(err)
+	}
+
+	if respBody["version"] != "1.0" {
+		t.Fatalf("expected version='1.0' in response body, got %v", respBody["version"])
+	}
+
+	// Check decision log
+	if logged == nil {
+		t.Fatal("expected decision log entry")
+	}
+	if logged.Custom == nil {
+		t.Fatal("expected Custom in decision log")
+	}
+
+	outgoing, ok := logged.Custom["response_metadata"].(map[string]any)
+	if !ok {
+		t.Fatalf("expected response_metadata in Custom, got %v", logged.Custom)
+	}
+	if outgoing["version"] != "1.0" {
+		t.Fatalf("expected version='1.0' in response_metadata, got %v", outgoing["version"])
+	}
+
+	// Request metadata should also be present
+	if _, ok := logged.Custom["request_metadata"].(map[string]any); !ok {
+		t.Fatal("expected request_metadata in Custom")
+	}
+}
+
 func evalReq(t testing.TB, path string, payload map[string]any, target string) *http.Request {
 	t.Helper()