From d271cdfff8b771402c8b3b4d37281b39f1c937f8 Mon Sep 17 00:00:00 2001
From: lif <1835304752@qq.com>
Date: Fri, 9 Jan 2026 06:02:23 +0800
Subject: [PATCH] rego: Add Data function to simplify adding data from map
 (#8166)

Add rego.Data function to allow setting data directly from a
map[string]any, providing a simpler alternative to using
Store(inmem.NewFromObject(data)). This improves the Go SDK API
by reducing boilerplate for the common case of using an in-memory
store with static data.

Fixes: #5961

Signed-off-by: majiayu000 <1835304752@qq.com>
---
 rego/rego.go         |  8 ++++++
 v1/rego/rego.go      | 10 +++++++
 v1/rego/rego_test.go | 68 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+)

diff --git a/rego/rego.go b/rego/rego.go
index 0727dae69a..c9caf9f8cb 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -408,6 +408,14 @@ func Store(s storage.Store) func(r *Rego) {
 	return v1.Store(s)
 }
 
+// Data returns an argument that sets the Rego data document. Data should be
+// a map representing the data document. This is a simpler alternative to
+// using Store with inmem.NewFromObject for cases where an in-memory store
+// with static data is sufficient.
+func Data(x map[string]any) func(r *Rego) {
+	return v1.Data(x)
+}
+
 // StoreReadAST returns an argument that sets whether the store should eagerly convert data to AST values.
 //
 // Only applicable when no store has been set on the Rego object through the Store option.
diff --git a/v1/rego/rego.go b/v1/rego/rego.go
index fb95836d1a..1e90cd091a 100644
--- a/v1/rego/rego.go
+++ b/v1/rego/rego.go
@@ -1080,6 +1080,16 @@ func Store(s storage.Store) func(r *Rego) {
 	}
 }
 
+// Data returns an argument that sets the Rego data document. Data should be
+// a map representing the data document. This is a simpler alternative to
+// using Store with inmem.NewFromObject for cases where an in-memory store
+// with static data is sufficient.
+func Data(x map[string]any) func(r *Rego) {
+	return func(r *Rego) {
+		r.store = inmem.NewFromObject(x)
+	}
+}
+
 // StoreReadAST returns an argument that sets whether the store should eagerly convert data to AST values.
 //
 // Only applicable when no store has been set on the Rego object through the Store option.
diff --git a/v1/rego/rego_test.go b/v1/rego/rego_test.go
index 5c9ffeb3b6..88e5b17828 100644
--- a/v1/rego/rego_test.go
+++ b/v1/rego/rego_test.go
@@ -3449,3 +3449,71 @@ result := test.module("policy.rego")
 		}
 	})
 }
+
+func TestRegoData(t *testing.T) {
+	ctx := t.Context()
+
+	r := New(
+		Query("data.x.y"),
+		Data(map[string]any{
+			"x": map[string]any{
+				"y": "hello",
+			},
+		}),
+	)
+
+	rs, err := r.Eval(ctx)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if len(rs) != 1 || len(rs[0].Expressions) != 1 {
+		t.Fatalf("Expected one result with one expression but got: %v", rs)
+	}
+
+	if rs[0].Expressions[0].Value != "hello" {
+		t.Fatalf("Expected 'hello' but got: %v", rs[0].Expressions[0].Value)
+	}
+}
+
+func TestRegoDataWithModule(t *testing.T) {
+	ctx := t.Context()
+
+	mod := `
+	package test
+	import rego.v1
+
+	result := data.users[input.user_id].role
+	`
+
+	r := New(
+		Query("data.test.result"),
+		Module("test.rego", mod),
+		Data(map[string]any{
+			"users": map[string]any{
+				"alice": map[string]any{
+					"role": "admin",
+				},
+				"bob": map[string]any{
+					"role": "viewer",
+				},
+			},
+		}),
+		Input(map[string]any{
+			"user_id": "alice",
+		}),
+	)
+
+	rs, err := r.Eval(ctx)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if len(rs) != 1 || len(rs[0].Expressions) != 1 {
+		t.Fatalf("Expected one result with one expression but got: %v", rs)
+	}
+
+	if rs[0].Expressions[0].Value != "admin" {
+		t.Fatalf("Expected 'admin' but got: %v", rs[0].Expressions[0].Value)
+	}
+}