From 578fa9e3cd9d50b06da0d38241e6cb89edb76f50 Mon Sep 17 00:00:00 2001 From: Sebastian Spaink Date: Mon, 10 Aug 2026 11:49:58 -0500 Subject: [PATCH] repl: Allow interactive ref head rule definitions Statements such as `a[0] := 1` or `p.q.r := 1` were rejected with "cannot assign to ref". The REPL refused to interpret refs of more than one term as rule heads, so those statements fell through to being compiled as a query body, where assigning to a ref isn't allowed. They are now interpreted as rule definitions. Refs rooted at data or input are excluded, as `data.foo.bar = 1` asks whether that document is 1 rather than defining a rule. Rules are identified by their head ref instead of Head.Name when unsetting them, which fixes two related problems: ref head rules (so far only definable with the `if` keyword) could not be unset at all, because their Head.Name is empty, and re-defining `a[0] := 1` doesn't drop unrelated keys of the same document, e.g. a[1]. As a consequence, `unset` accepts a ref, e.g. `unset a[0]` or `unset p.q.r`, and removes all rules below it. ParsePartialObjectDocRuleFromEqExpr didn't mark the `true` body it generates as generated, so rules parsed from these statements tripped the rego-v1 check requiring `if` before a rule body. The module parser sets that flag itself after calling ParseRuleFromBody, which is why this only surfaced for direct callers like the REPL. Fixes: #5498 Signed-off-by: Sebastian Spaink --- v1/ast/parser_ext.go | 9 ++- v1/ast/parser_test.go | 30 +++++++ v1/repl/repl.go | 122 +++++++++++++++++++++++------ v1/repl/repl_test.go | 177 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 309 insertions(+), 29 deletions(-) diff --git a/v1/ast/parser_ext.go b/v1/ast/parser_ext.go index 8b63182c0f9..557f7588afa 100644 --- a/v1/ast/parser_ext.go +++ b/v1/ast/parser_ext.go @@ -346,10 +346,11 @@ func ParsePartialObjectDocRuleFromEqExpr(module *Module, lhs, rhs *Term) (*Rule, body := NewBody(NewExpr(BooleanTerm(true).SetLocation(rhs.Location)).SetLocation(rhs.Location)) rule := &Rule{ - Location: rhs.Location, - Head: head, - Body: body, - Module: module, + Location: rhs.Location, + Head: head, + Body: body, + Module: module, + generatedBody: true, } return rule, nil diff --git a/v1/ast/parser_test.go b/v1/ast/parser_test.go index a31a8d9b407..b190d2788a6 100644 --- a/v1/ast/parser_test.go +++ b/v1/ast/parser_test.go @@ -5773,6 +5773,36 @@ data = {"bar": 2}` } } +func TestRuleFromExprGeneratedBody(t *testing.T) { + // Rules parsed from a single expression have a generated body, so no `if` + // keyword is required of them under rego-v1. This is what allows rules to + // be defined interactively in the REPL. + tests := []string{ + `x := 1`, + `x = 1`, + `a[0] := 1`, + `a["foo"] = "bar"`, + `p.q.r := 1`, + } + + mod := MustParseModule("package test") + + for _, tc := range tests { + t.Run(tc, func(t *testing.T) { + rule, err := ParseRuleFromExpr(mod, MustParseBody(tc)[0]) + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if !rule.generatedBody { + t.Fatal("Expected rule to have generated body") + } + if errs := CheckRegoV1(rule); len(errs) > 0 { + t.Fatalf("Unexpected errors: %v", errs) + } + }) + } +} + func TestWildcards(t *testing.T) { assertParseOneTerm(t, "ref", "a.b[_].c[_]", RefTerm( diff --git a/v1/repl/repl.go b/v1/repl/repl.go index 5b98cd1e6de..2b6857c1d60 100644 --- a/v1/repl/repl.go +++ b/v1/repl/repl.go @@ -731,7 +731,7 @@ func (r *REPL) cmdUnknown(s []string) error { func (r *REPL) cmdUnset(ctx context.Context, args []string) error { if len(args) != 1 { - return newBadArgsErr("unset : expects exactly one argument") + return newBadArgsErr("unset : expects exactly one argument") } term, err := ast.ParseTerm(args[0]) @@ -739,17 +739,12 @@ func (r *REPL) cmdUnset(ctx context.Context, args []string) error { return newBadArgsErr("argument must identify a rule") } - v, ok := term.Value.(ast.Var) - + ref, ok := ruleRefFromTerm(term) if !ok { - ref, ok := term.Value.(ast.Ref) - if !ok || !ast.RootDocumentNames.Contains(ref[0]) { - return newBadArgsErr("arguments must identify a rule") - } - v = ref[0].Value.(ast.Var) + return newBadArgsErr("arguments must identify a rule") } - unset, err := r.unsetRule(ctx, v) + unset, err := r.unsetRule(ctx, ref) if err != nil { return err } else if !unset { @@ -779,7 +774,7 @@ func (r *REPL) cmdUnsetPackage(ctx context.Context, args []string) error { return nil } -func (r *REPL) unsetRule(ctx context.Context, name ast.Var) (bool, error) { +func (r *REPL) unsetRule(ctx context.Context, ref ast.Ref) (bool, error) { if r.currentModuleID == "" { return false, nil } @@ -787,9 +782,9 @@ func (r *REPL) unsetRule(ctx context.Context, name ast.Var) (bool, error) { mod := r.modules[r.currentModuleID] rules := []*ast.Rule{} - for _, r := range mod.Rules { - if r.Head.Name != name { - rules = append(rules, r) + for _, rule := range mod.Rules { + if !refsOverlap(rule.Head.Ref(), ref) { + rules = append(rules, rule) } } @@ -807,6 +802,45 @@ func (r *REPL) unsetRule(ctx context.Context, name ast.Var) (bool, error) { return true, nil } +// ruleRefFromTerm returns the rule head ref identified by term: the terms "p", +// "p.q.r" and `p["q"]` identify the rule head refs p, p.q.r and p.q. Refs +// rooted at a root document are only identified by their first term, e.g. +// "input" identifies the rule defined by `input = {...}`. +func ruleRefFromTerm(term *ast.Term) (ast.Ref, bool) { + switch v := term.Value.(type) { + case ast.Var: + return ast.Ref{term}, true + case ast.Ref: + if _, ok := v[0].Value.(ast.Var); !ok { + return nil, false + } + if ast.RootDocumentNames.Contains(v[0]) { + return v[:1], true + } + // Ref.IsGround ignores the leading var, so this only rejects refs with + // non-ground terms after the head, e.g. "a[x]". + if !v.IsGround() { + return nil, false + } + return v, true + } + return nil, false +} + +// refsOverlap returns true if one of the refs is a prefix of, or equal to, the +// other. Rule head refs that overlap in this way define (parts of) the same +// document: "user" overlaps "user.role", so unsetting "user" removes +// `user.role := "admin"`, while defining `user.role := "admin"` only replaces +// previous definitions of user.role, leaving user.email in place. +func refsOverlap(a, b ast.Ref) bool { + for i := range min(len(a), len(b)) { + if a[i].Value.Compare(b[i].Value) != 0 { + return false + } + } + return true +} + func (r *REPL) unsetPackage(_ context.Context, pkg *ast.Package) (bool, error) { path := pkg.Path.String() _, ok := r.modules[path] @@ -895,7 +929,7 @@ func (r *REPL) compileRule(ctx context.Context, rule *ast.Rule) error { if rule.Head.Assign { var err error - unset, err = r.unsetRule(ctx, rule.Head.Name) + unset, err = r.unsetRule(ctx, rule.Head.Ref()) if err != nil { return err } @@ -1295,6 +1329,12 @@ func (r *REPL) evalPackage(p *ast.Package) error { // > a // 1 // +// The left hand side may be a ref, defining part of a document: +// +// > user["role"] := "admin" +// > user +// {"role": "admin"} +// // If the expression is a = statement, then an additional check on the left // hand side occurs. For example: // @@ -1324,8 +1364,11 @@ func (r *REPL) interpretAsRule(ctx context.Context, compiler *ast.Compiler, body if rule == nil || err != nil { return false, err } - // TODO(sr): support interactive ref head rule definitions - if len(rule.Head.Ref()) > 1 { + + // Statements about a root document are queries, never rule definitions: + // "data.foo.bar = 1" asks if data.foo.bar is 1. The single-term case is + // excluded so that `input = {...}` keeps defining a rule. + if ref := rule.Head.Ref(); len(ref) > 1 && ast.RootDocumentNames.Contains(ref[0]) { return false, nil } @@ -1600,7 +1643,7 @@ var extra = [...]commandDesc{ var builtin = [...]commandDesc{ {"show", []string{""}, "show active module definition"}, {"show debug", []string{""}, "show REPL settings"}, - {"unset", []string{""}, "unset rules in currently active module"}, + {"unset", []string{""}, "unset rules in currently active module"}, {"unset-package", []string{""}, "unset packages in currently active module"}, {"json", []string{}, "set output format to JSON"}, {"pretty", []string{}, "set output format to pretty"}, @@ -1662,35 +1705,64 @@ func dumpStorage(ctx context.Context, store storage.Store, txn storage.Transacti return e.Encode(data) } +// isGlobalInModule returns true if term refers to an import, or to a document +// that the module already defines rules for. Statements about such documents +// are evaluated as queries rather than interpreted as rule definitions. func isGlobalInModule(compiler *ast.Compiler, module *ast.Module, term *ast.Term) bool { - var name ast.Var + var ref ast.Ref - if ast.RootDocumentRefs.Contains(term) { - name = term.Value.(ast.Ref)[0].Value.(ast.Var) - } else if v, ok := term.Value.(ast.Var); ok { - name = v - } else { + switch v := term.Value.(type) { + case ast.Var: + ref = ast.Ref{term} + case ast.Ref: + if _, ok := v[0].Value.(ast.Var); !ok { + return false + } + ref = v + default: return false } + name := ref[0].Value.(ast.Var) + for _, imp := range module.Imports { if imp.Name().Compare(name) == 0 { return true } } - path := module.Package.Path.Copy().Append(ast.StringTerm(string(name))) + // Only the ground prefix of the ref can be looked up in the rule tree, e.g. + // "a[i]" is looked up as "a". The rule tree keys the leading var of the ref + // as a string. + prefix := ref.GroundPrefix() + path := make(ast.Ref, 0, len(prefix)) + path = append(path, ast.StringTerm(string(name))) + path = append(path, prefix[1:]...) + node := compiler.RuleTree + for _, elem := range module.Package.Path { + node = node.Child(elem.Value) + if node == nil { + return false + } + } + for _, elem := range path { node = node.Child(elem.Value) if node == nil { return false } + if len(node.Values) > 0 { + return true + } } - return len(node.Values) > 0 + // The whole prefix resolved to a node without rules of its own: the ref + // refers to an existing document only if the dropped suffix can be + // satisfied by rules below that node, e.g. "a[i]" when a[0] is defined. + return len(prefix) < len(ref) } func printHelp(output io.Writer, initPrompt string, report [][2]string) { diff --git a/v1/repl/repl_test.go b/v1/repl/repl_test.go index 361a3bcae6b..36ea3949ba1 100644 --- a/v1/repl/repl_test.go +++ b/v1/repl/repl_test.go @@ -1222,6 +1222,183 @@ func TestOneShotRefHeadRulePrinted(t *testing.T) { expectOutput(t, buffer.String(), "Rule 'foo.bar.baz' defined in package repl. Type 'show' to see rules.\n") } +// Ref head rules can be defined interactively, i.e. without the `if` keyword +// and a rule body, see https://github.com/open-policy-agent/opa/issues/5498 +func TestOneShotRefHeadRuleDefinition(t *testing.T) { + tests := []struct { + note string + stmts []string + exp []string + }{ + { + note: "number key", + stmts: []string{`a[0] := 1`, `a`}, + exp: []string{ + "Rule 'a[0]' defined in package repl. Type 'show' to see rules.\n", + "{\n \"0\": 1\n}\n", + }, + }, + { + note: "string key", + stmts: []string{`a["foo"] := "bar"`, `a`}, + exp: []string{ + "Rule 'a.foo' defined in package repl. Type 'show' to see rules.\n", + "{\n \"foo\": \"bar\"\n}\n", + }, + }, + { + note: "distinct keys are kept", + stmts: []string{`a[0] := 1`, `a[1] := 2`, `a`}, + exp: []string{ + "Rule 'a[0]' defined in package repl. Type 'show' to see rules.\n", + "Rule 'a[1]' defined in package repl. Type 'show' to see rules.\n", + "{\n \"0\": 1,\n \"1\": 2\n}\n", + }, + }, + { + note: "same key is re-defined", + stmts: []string{`a[0] := 1`, `a[1] := 2`, `a[0] := 3`, `a`}, + exp: []string{ + "Rule 'a[0]' defined in package repl. Type 'show' to see rules.\n", + "Rule 'a[1]' defined in package repl. Type 'show' to see rules.\n", + "Rule 'a[0]' re-defined in package repl. Type 'show' to see rules.\n", + "{\n \"0\": 3,\n \"1\": 2\n}\n", + }, + }, + { + note: "complete rule replaces keys", + stmts: []string{`a[0] := 1`, `a := 2`, `a`}, + exp: []string{ + "Rule 'a[0]' defined in package repl. Type 'show' to see rules.\n", + "Rule 'a' re-defined in package repl. Type 'show' to see rules.\n", + "2\n", + }, + }, + { + note: "dotted ref", + stmts: []string{`p.q.r := 1`, `p.q.s := 2`, `p`}, + exp: []string{ + "Rule 'p.q.r' defined in package repl. Type 'show' to see rules.\n", + "Rule 'p.q.s' defined in package repl. Type 'show' to see rules.\n", + "{\n \"q\": {\n \"r\": 1,\n \"s\": 2\n }\n}\n", + }, + }, + { + note: "assignment to var key is unsafe", + stmts: []string{`a[i] := 1`}, + exp: []string{""}, + }, + { + note: "eq statement defines rule", + stmts: []string{`p.q.r = 1`, `p.q.r`}, + exp: []string{ + "Rule 'p.q.r' defined in package repl. Type 'show' to see rules.\n", + "1\n", + }, + }, + { + // The rule isn't re-defined, the statement is a query about the + // existing document. + note: "eq statement about defined rule is a query", + stmts: []string{`p.q.r := 1`, `p.q.r = 1`, `p.q.r = 2`, `p.q[i] = 1`}, + exp: []string{ + "Rule 'p.q.r' defined in package repl. Type 'show' to see rules.\n", + "true\n", + "undefined\n", + "┌─────┐\n│ i │\n├─────┤\n│ \"r\" │\n└─────┘\n", + }, + }, + { + note: "data ref is a query", + stmts: []string{`data.foo.bar = 1`, `show`}, + exp: []string{ + "undefined\n", + "no rules defined\n", + }, + }, + { + note: "input ref is a query", + stmts: []string{`input.foo.bar = 1`, `show`}, + exp: []string{ + "undefined\n", + "no rules defined\n", + }, + }, + } + + for _, tc := range tests { + t.Run(tc.note, func(t *testing.T) { + ctx := t.Context() + var buffer bytes.Buffer + repl := newRepl(inmem.New(), &buffer) + + for i, stmt := range tc.stmts { + buffer.Reset() + err := repl.OneShot(ctx, stmt) + if tc.exp[i] == "" { + if err == nil { + t.Fatalf("%q: expected error but got output: %q", stmt, buffer.String()) + } + continue + } + if err != nil { + t.Fatalf("%q: unexpected error: %v", stmt, err) + } + if act := buffer.String(); act != tc.exp[i] { + t.Fatalf("%q: expected output %q but got %q", stmt, tc.exp[i], act) + } + } + }) + } +} + +// Ref head rules are identified by their ref, e.g. "unset a[0]" and +// "unset foo.bar.baz". +func TestUnsetRefHeadRule(t *testing.T) { + ctx := t.Context() + var buffer bytes.Buffer + repl := newRepl(inmem.New(), &buffer) + + for _, stmt := range []string{`a[0] := 1`, `a[1] := 2`, `foo.bar.baz if true`} { + if err := repl.OneShot(ctx, stmt); err != nil { + t.Fatalf("%q: unexpected error: %v", stmt, err) + } + } + + // Unsetting a key leaves the other keys of the document in place. + buffer.Reset() + if err := repl.OneShot(ctx, `unset a[0]`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if err := repl.OneShot(ctx, `a`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + expectOutput(t, buffer.String(), "{\n \"1\": 2\n}\n") + + // Unsetting a prefix of the ref removes all rules under it. + buffer.Reset() + if err := repl.OneShot(ctx, `unset a`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if err := repl.OneShot(ctx, `unset foo.bar.baz`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + if err := repl.OneShot(ctx, `show`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + expectOutput(t, buffer.String(), "package repl\n") + + buffer.Reset() + if err := repl.OneShot(ctx, `unset foo.bar.baz`); err != nil { + t.Fatalf("Unexpected error: %v", err) + } + expectOutput(t, buffer.String(), "warning: no matching rules in current module\n") + + if err := repl.OneShot(ctx, `unset a[x]`); err == nil { + t.Fatal("Expected error for non-ground ref") + } +} + func TestOneShotBufferedExpr(t *testing.T) { ctx := t.Context() store := newTestStore()