rego: add ResultValue[T]() helper method

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

v1/rego/resultset.go

@@ -79,12 +79,24 @@ func (ev *ExpressionValue) String() string {
 // return `true` for a query like `data.authz.allow = x`, which always has result
 // set element with value true, but could also have a binding `x: false`.
 func (rs ResultSet) Allowed() bool {
+	x, _ := ResultValue[bool](rs)
+	return x
+}
+
+// ResultValue is a helper function that'll return a value of type T if all of
+// these conditions hold:
+// - the result set only has one element
+// - there is only one expression in the result set's only element
+// - that expression has type T
+// - there are no bindings.
+func ResultValue[T any](rs ResultSet) (T, bool) {
+	var zero T
 	if len(rs) == 1 && len(rs[0].Bindings) == 0 {
 		if exprs := rs[0].Expressions; len(exprs) == 1 {
-			if b, ok := exprs[0].Value.(bool); ok {
-				return b
+			if v, ok := exprs[0].Value.(T); ok {
+				return v, true
 			}
 		}
 	}
-	return false
+	return zero, false
 }

v1/rego/resultset_test.go

@@ -16,8 +16,7 @@ func TestResultSetAllowed(t *testing.T) {
 		{
 			note: "simplest true",
 			module: `package authz
-import rego.v1
-allow if { true }
+allow if true
 `,
 			query:    "data.authz.allow",
 			expected: true,
@@ -25,7 +24,7 @@ allow if { true }
 		{
 			note: "simplest false",
 			module: `package authz
-default allow = false
+default allow := false
 `,
 			query:    "data.authz.allow",
 			expected: false,
@@ -33,8 +32,7 @@ default allow = false
 		{
 			note: "true value + bindings",
 			module: `package authz
-import rego.v1
-allow if { true }
+allow if true
 `,
 			query:    "data.authz.allow = x",
 			expected: false,
@@ -42,8 +40,7 @@ allow if { true }
 		{
 			note: "object response, bound to var in query",
 			module: `package authz
-import rego.v1
-resp = { "allow": true } if { true }
+resp := {"allow": true}
 `,
 			query:    "data.authz.resp = x",
 			expected: false,
@@ -51,8 +48,7 @@ resp = { "allow": true } if { true }
 		{
 			note: "object response, treated as false",
 			module: `package authz
-import rego.v1
-resp = { "allow": true } if { true }
+resp := {"allow": true}
 `,
 			query:    "data.authz.resp",
 			expected: false,
@@ -75,3 +71,158 @@ resp = { "allow": true } if { true }
 		})
 	}
 }
+
+func TestResultValue(t *testing.T) {
+	t.Run("bool", func(t *testing.T) {
+		tests := []struct {
+			note          string
+			module        string
+			query         string
+			expectedValue bool
+			expectedOk    bool
+		}{
+			{
+				note: "true value",
+				module: `package authz
+allow if true
+`,
+				query:         "data.authz.allow",
+				expectedValue: true,
+				expectedOk:    true,
+			},
+			{
+				note: "false value",
+				module: `package authz
+default allow := false
+`,
+				query:         "data.authz.allow",
+				expectedValue: false,
+				expectedOk:    true,
+			},
+			{
+				note: "value with bindings",
+				module: `package authz
+allow if true
+`,
+				query:         "data.authz.allow = x",
+				expectedValue: false,
+				expectedOk:    false,
+			},
+		}
+
+		for _, tc := range tests {
+			t.Run(tc.note, func(t *testing.T) {
+				r := rego.New(
+					rego.Query(tc.query),
+					rego.Module("", tc.module),
+				)
+				rs, err := r.Eval(t.Context())
+				if err != nil {
+					t.Fatal(err)
+				}
+				val, ok := rego.ResultValue[bool](rs)
+				if exp, act := tc.expectedOk, ok; exp != act {
+					t.Errorf("expected ok=%v, got ok=%v", exp, act)
+				}
+				if ok && tc.expectedOk {
+					if exp, act := tc.expectedValue, val; exp != act {
+						t.Errorf("expected value=%v, got value=%v", exp, act)
+					}
+				}
+			})
+		}
+	})
+
+	t.Run("string", func(t *testing.T) {
+		tests := []struct {
+			note          string
+			module        string
+			query         string
+			expectedValue string
+			expectedOk    bool
+		}{
+			{
+				note: "string value",
+				module: `package authz
+message := "hello world"
+`,
+				query:         "data.authz.message",
+				expectedValue: "hello world",
+				expectedOk:    true,
+			},
+			{
+				note: "empty string",
+				module: `package authz
+message := ""
+`,
+				query:         "data.authz.message",
+				expectedValue: "",
+				expectedOk:    true,
+			},
+		}
+
+		for _, tc := range tests {
+			t.Run(tc.note, func(t *testing.T) {
+				r := rego.New(
+					rego.Query(tc.query),
+					rego.Module("", tc.module),
+				)
+				rs, err := r.Eval(t.Context())
+				if err != nil {
+					t.Fatal(err)
+				}
+				val, ok := rego.ResultValue[string](rs)
+				if exp, act := tc.expectedOk, ok; exp != act {
+					t.Errorf("expected ok=%v, got ok=%v", exp, act)
+				}
+				if ok && tc.expectedOk {
+					if exp, act := tc.expectedValue, val; exp != act {
+						t.Errorf("expected value=%q, got value=%q", exp, act)
+					}
+				}
+			})
+		}
+	})
+
+	t.Run("wrong type", func(t *testing.T) {
+		module := `package authz
+message := "hello world"
+`
+		r := rego.New(
+			rego.Query("data.authz.message"),
+			rego.Module("", module),
+		)
+		rs, err := r.Eval(t.Context())
+		if err != nil {
+			t.Fatal(err)
+		}
+		val, ok := rego.ResultValue[int](rs)
+		if ok {
+			t.Errorf("expected ok=false for wrong type conversion, got ok=true with value=%v", val)
+		}
+	})
+
+	t.Run("object value", func(t *testing.T) {
+		module := `package authz
+resp := {"allow": true, "user": "alice"}
+`
+		r := rego.New(
+			rego.Query("data.authz.resp"),
+			rego.Module("", module),
+		)
+		rs, err := r.Eval(t.Context())
+		if err != nil {
+			t.Fatal(err)
+		}
+		val, ok := rego.ResultValue[map[string]any](rs)
+		if !ok {
+			t.Fatal("expected ok=true for map type")
+		}
+		if exp, act := true, val["allow"]; exp != act {
+			t.Errorf("expected allow=%v, got allow=%v", exp, act)
+		}
+		if exp, act := "alice", val["user"]; exp != act {
+			t.Errorf("expected user=%v, got user=%v", exp, act)
+		}
+	})
+}