Compile API: switch to compile annotation key (#7936)

* ast: add compile annotation key

This used to be

```
  custom:
    unknowns: [ ... ]
    mask_rule: ...
```

and now becomes
```
  compile:
    unknowns: [ ... ]
    mask_rule: ...
```

* server: adapt compile handler annotations processing

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
This commit is contained in:
Stephan Renatus
2025-09-24 10:53:17 +02:00
committed by GitHub
parent 9c1cf16d4b
commit 85a0e2a28c
15 changed files with 213 additions and 80 deletions
+1 -1
View File
@@ -1664,7 +1664,7 @@ package filters
# METADATA
# scope: document
# custom:
# compile:
# unknowns: [input.fruits]
include if input.fruits.name == input.favorite
+1 -1
View File
@@ -566,7 +566,7 @@ func TestPrometheusMetrics(t *testing.T) {
policy := `package filters
# METADATA
# scope: document
# custom:
# compile:
# unknowns: [input.fruits]
include if input.fruits.name == "banana"
`
+1 -1
View File
@@ -37,7 +37,7 @@ func TestDecisionLogsCompileAPIResult(t *testing.T) {
package filters
# METADATA
# custom:
# compile:
# unknowns: [input.fruits]
# mask_rule: data.filters.mask
include if input.fruits.name in input.favorites
+17 -43
View File
@@ -7,7 +7,6 @@ import (
"encoding/json"
"fmt"
"slices"
"strings"
"github.com/open-policy-agent/opa/internal/levenshtein"
"github.com/open-policy-agent/opa/internal/ucast"
@@ -15,8 +14,7 @@ import (
)
const (
invalidUnknownCode = "invalid_unknown"
invalidMaskRuleCode = "invalid_mask_rule"
invalidUnknownCode = "invalid_unknown"
)
type UCASTNode struct {
@@ -58,7 +56,7 @@ func QueriesToSQL(queries []ast.Body, mappings map[string]any, dialect string) (
return sql, nil
}
func ExtractUnknownsFromAnnotations(comp *ast.Compiler, ref ast.Ref) ([]*ast.Term, []*ast.Error) {
func ExtractUnknownsFromAnnotations(comp *ast.Compiler, ref ast.Ref) ([]ast.Ref, []*ast.Error) {
// find ast.Rule for ref
rules := comp.GetRulesExact(ref)
if len(rules) == 0 {
@@ -68,36 +66,22 @@ func ExtractUnknownsFromAnnotations(comp *ast.Compiler, ref ast.Ref) ([]*ast.Ter
return unknownsFromAnnotationsSet(comp.GetAnnotationSet(), rule)
}
func unknownsFromAnnotationsSet(as *ast.AnnotationSet, rule *ast.Rule) ([]*ast.Term, []*ast.Error) {
func unknownsFromAnnotationsSet(as *ast.AnnotationSet, rule *ast.Rule) ([]ast.Ref, []*ast.Error) {
if as == nil {
return nil, nil
}
var unknowns []*ast.Term
var unknowns []ast.Ref
var errs []*ast.Error
for _, ar := range as.Chain(rule) {
ann := ar.Annotations
if ann == nil {
if ann == nil || ann.Compile == nil {
continue
}
unk, ok := ann.Custom["unknowns"]
if !ok {
continue
}
unkArray, ok := unk.([]any)
if !ok {
continue
}
for _, u := range unkArray {
s, ok := u.(string)
if !ok {
continue
}
ref, err := ast.ParseRef(s)
if err != nil {
errs = append(errs, ast.NewError(invalidUnknownCode, ann.Loc(), "unknowns must be valid refs: %s", s))
} else if ref.HasPrefix(ast.DefaultRootRef) || ref.HasPrefix(ast.InputRootRef) {
unknowns = append(unknowns, ast.NewTerm(ref))
unkArray := ann.Compile.Unknowns
for _, ref := range unkArray {
if ref.HasPrefix(ast.DefaultRootRef) || ref.HasPrefix(ast.InputRootRef) {
unknowns = append(unknowns, ref)
} else {
errs = append(errs, ast.NewError(invalidUnknownCode, ann.Loc(), "unknowns must be prefixed with `input` or `data`: %v", ref))
}
@@ -114,35 +98,25 @@ func ExtractMaskRuleRefFromAnnotations(comp *ast.Compiler, ref ast.Ref) (ast.Ref
return nil, nil
}
rule := rules[0] // rule scope doesn't make sense here, so it doesn't matter which rule we use
return maskRuleFromAnnotationsSet(comp.GetAnnotationSet(), comp, rule)
return maskRuleFromAnnotationsSet(comp.GetAnnotationSet(), rule)
}
func maskRuleFromAnnotationsSet(as *ast.AnnotationSet, comp *ast.Compiler, rule *ast.Rule) (ast.Ref, *ast.Error) {
func maskRuleFromAnnotationsSet(as *ast.AnnotationSet, rule *ast.Rule) (ast.Ref, *ast.Error) {
if as == nil {
return nil, nil
}
for _, ar := range as.Chain(rule) {
ann := ar.Annotations
if ann == nil {
if ann == nil || ann.Compile == nil {
continue
}
// If the mask_rule key is present, validate and parse it.
if maskRule, ok := ann.Custom["mask_rule"]; ok {
if s, ok := maskRule.(string); ok {
maskPath := s
if !strings.HasPrefix(s, "data.") {
// If the mask_rule is not a data ref try adding package prefix.
maskPath = rule.Module.Package.Path.String() + "." + s
}
maskRuleRef, err := ast.ParseRef(maskPath)
if err != nil {
hint := FuzzyRuleNameMatchHint(comp, s)
return nil, ast.NewError(invalidMaskRuleCode, ann.Loc(), "mask_rule was not a valid ref: %s", hint)
}
return maskRuleRef, nil
if maskRule := ann.Compile.MaskRule; maskRule != nil {
if !maskRule.HasPrefix(ast.DefaultRootRef) {
// If the mask_rule is not a data ref, add package prefix.
maskRule = rule.Module.Package.Path.Extend(maskRule)
}
return nil, ast.NewError(invalidMaskRuleCode, ann.Loc(), "mask_rule must be a valid ref string: %v", maskRule)
return maskRule, nil
}
}
+50
View File
@@ -34,6 +34,7 @@ type (
RelatedResources []*RelatedResourceAnnotation `json:"related_resources,omitempty"`
Authors []*AuthorAnnotation `json:"authors,omitempty"`
Schemas []*SchemaAnnotation `json:"schemas,omitempty"`
Compile *CompileAnnotation `json:"compile,omitempty"`
Custom map[string]any `json:"custom,omitempty"`
Location *Location `json:"location,omitempty"`
@@ -48,6 +49,11 @@ type (
Definition *any `json:"definition,omitempty"`
}
CompileAnnotation struct {
Unknowns []Ref `json:"unknowns,omitempty"`
MaskRule Ref `json:"mask_rule,omitempty"` // NOTE: This doesn't need to start with "data.package", it can be relative
}
AuthorAnnotation struct {
Name string `json:"name"`
Email string `json:"email,omitempty"`
@@ -151,6 +157,10 @@ func (a *Annotations) Compare(other *Annotations) int {
return cmp
}
if cmp := a.Compile.Compare(other.Compile); cmp != 0 {
return cmp
}
if a.Entrypoint != other.Entrypoint {
if a.Entrypoint {
return 1
@@ -403,6 +413,8 @@ func (a *Annotations) Copy(node Node) *Annotations {
cpy.Schemas[i] = a.Schemas[i].Copy()
}
cpy.Compile = a.Compile.Copy()
if a.Custom != nil {
cpy.Custom = deepcopy.Map(a.Custom)
}
@@ -716,6 +728,44 @@ func (s *SchemaAnnotation) String() string {
return string(bs)
}
// Copy returns a deep copy of s.
func (c *CompileAnnotation) Copy() *CompileAnnotation {
if c == nil {
return nil
}
cpy := *c
for i := range c.Unknowns {
cpy.Unknowns[i] = c.Unknowns[i].Copy()
}
return &cpy
}
// Compare returns an integer indicating if s is less than, equal to, or greater
// than other.
func (c *CompileAnnotation) Compare(other *CompileAnnotation) int {
switch {
case c == nil && other == nil:
return 0
case c != nil && other == nil:
return 1
case c == nil && other != nil:
return -1
}
if cmp := slices.CompareFunc(c.Unknowns, other.Unknowns,
func(x, y Ref) int {
return x.Compare(y)
}); cmp != 0 {
return cmp
}
return c.MaskRule.Compare(other.MaskRule)
}
func (c *CompileAnnotation) String() string {
bs, _ := json.Marshal(c)
return string(bs)
}
func newAnnotationSet() *AnnotationSet {
return &AnnotationSet{
byRule: map[*Rule][]*Annotations{},
+35
View File
@@ -2570,6 +2570,7 @@ type rawAnnotation struct {
RelatedResources []any `yaml:"related_resources"`
Authors []any `yaml:"authors"`
Schemas []map[string]any `yaml:"schemas"`
Compile map[string]any `yaml:"compile"`
Custom map[string]any `yaml:"custom"`
}
@@ -2637,6 +2638,40 @@ func (b *metadataParser) Parse() (*Annotations, error) {
result.RelatedResources = append(result.RelatedResources, rr)
}
if raw.Compile != nil {
result.Compile = &CompileAnnotation{}
if unknowns, ok := raw.Compile["unknowns"]; ok {
if unknowns, ok := unknowns.([]any); ok {
result.Compile.Unknowns = make([]Ref, len(unknowns))
for i := range unknowns {
if unknown, ok := unknowns[i].(string); ok {
ref, err := ParseRef(unknown)
if err != nil {
return nil, fmt.Errorf("invalid unknowns element %q: %w", unknown, err)
}
result.Compile.Unknowns[i] = ref
}
}
}
}
if mask, ok := raw.Compile["mask_rule"]; ok {
if mask, ok := mask.(string); ok {
maskTerm, err := ParseTerm(mask)
if err != nil {
return nil, fmt.Errorf("invalid mask_rule annotation %q: %w", mask, err)
}
switch v := maskTerm.Value.(type) {
case Var, String:
result.Compile.MaskRule = Ref{maskTerm}
case Ref:
result.Compile.MaskRule = v
default:
return nil, fmt.Errorf("invalid mask_rule annotation type %q: %[1]T", mask)
}
}
}
}
for _, pair := range raw.Schemas {
k, v := unwrapPair(pair)
+72
View File
@@ -7008,6 +7008,78 @@ p if { input = "str" }`,
},
},
},
{
note: "compile annotation, short mask_rule",
module: `
package opa.examples
# METADATA
# scope: document
# compile:
# unknowns:
# - input.fruits
# mask_rule: mask
include if input.fruits.name == "banana"
mask.fruits.owner.replace.value := "___"
`,
expNumComments: 6,
expAnnotations: []*Annotations{
{
Scope: annotationScopeDocument,
Compile: &CompileAnnotation{
Unknowns: []Ref{MustParseRef("input.fruits")},
MaskRule: EmptyRef().Append(VarTerm("mask")),
},
},
},
},
{
note: "compile annotation, full mask_rule",
module: `
package opa.examples
# METADATA
# scope: document
# compile:
# unknowns:
# - input.fruits
# mask_rule: data.filtering.mask
include if input.fruits.name == "banana"
mask.fruits.owner.replace.value := "___"
`,
expNumComments: 6,
expAnnotations: []*Annotations{
{
Scope: annotationScopeDocument,
Compile: &CompileAnnotation{
Unknowns: []Ref{MustParseRef("input.fruits")},
MaskRule: MustParseRef("data.filtering.mask"),
},
},
},
},
{
note: "compile annotation, no mask_rule",
module: `
package opa.examples
# METADATA
# scope: document
# compile:
# unknowns:
# - input.fruits
include if input.fruits.name == "banana"
`,
expNumComments: 5,
expAnnotations: []*Annotations{
{
Scope: annotationScopeDocument,
Compile: &CompileAnnotation{
Unknowns: []Ref{MustParseRef("input.fruits")},
},
},
},
},
}
for _, tc := range tests {
+1 -1
View File
@@ -26,7 +26,7 @@ type Compile struct {
targets []string
dialects []string
maskRule ast.Ref
unknowns []*ast.Term
unknowns []*ast.Term // ast.Ref would be slightly more on-the-spot, but we follow what is done in v1/rego to minimise surprises.
query ast.Body
mappings map[string]any
metrics metrics.Metrics
+11 -8
View File
@@ -93,7 +93,7 @@ type CompileFiltersRequestV1 struct {
type compileFiltersRequest struct {
Query ast.Body
Input ast.Value
Unknowns []*ast.Term
Unknowns []ast.Ref
Options compileFiltersRequestOptions
}
@@ -222,11 +222,15 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
targetOption = append(targetOption, rego_compile.Target(target, dialect))
}
unks := make([]*ast.Term, len(unknowns))
for i := range unknowns {
unks[i] = ast.NewTerm(unknowns[i])
}
m.Timer(timerPrepPartial).Start()
// NB(sr): just cache preparedCompile by path?
preparedCompile, err := rego_compile.New(
append(targetOption,
rego_compile.ParsedUnknowns(unknowns...),
rego_compile.ParsedUnknowns(unks...),
rego_compile.ParsedQuery(request.Query),
rego_compile.Metrics(m),
rego_compile.Mappings(orig.Options.Mappings),
@@ -356,7 +360,7 @@ func (s *Server) v1CompileFilters(w http.ResponseWriter, r *http.Request) {
}
}
func (s *Server) compileFiltersUnknowns(m metrics.Metrics, comp *ast.Compiler, path string, query ast.Body) ([]*ast.Term, []*ast.Error) {
func (s *Server) compileFiltersUnknowns(m metrics.Metrics, comp *ast.Compiler, path string, query ast.Body) ([]ast.Ref, []*ast.Error) {
key := path
unknowns, ok := s.compileUnknownsCache.Get(key)
if ok {
@@ -374,11 +378,10 @@ func (s *Server) compileFiltersUnknowns(m metrics.Metrics, comp *ast.Compiler, p
if !ok {
return nil, nil
}
parsedUnknowns, errs := compile.ExtractUnknownsFromAnnotations(comp, queryRef)
unknowns, errs := compile.ExtractUnknownsFromAnnotations(comp, queryRef)
if errs != nil {
return nil, errs
}
unknowns = parsedUnknowns
m.Timer(timerExtractAnnotationsUnknowns).Stop()
s.compileUnknownsCache.Add(key, unknowns)
@@ -452,11 +455,11 @@ func readInputCompileFiltersV1(comp *ast.Compiler, reqBytes []byte, urlPath stri
}
}
var unknowns []*ast.Term
var unknowns []ast.Ref
if request.Unknowns != nil {
unknowns = make([]*ast.Term, len(*request.Unknowns))
unknowns = make([]ast.Ref, len(*request.Unknowns))
for i, s := range *request.Unknowns {
unknowns[i], err = ast.ParseTerm(s)
unknowns[i], err = ast.ParseRef(s)
if err != nil {
return nil, nil, types.NewErrorV1(types.CodeInvalidParameter, "error(s) occurred while parsing unknowns: %v", err)
}
+13 -13
View File
@@ -80,7 +80,7 @@ func TestPostPartialChecks(t *testing.T) {
rego: `
# METADATA
# scope: document
# custom:
# compile:
# unknowns:
# - input.fruits
include if input.fruits.colour == input.colour
@@ -95,7 +95,7 @@ _use_metadata := rego.metadata.chain()`,
rego: `
# METADATA
# scope: package
# custom:
# compile:
# unknowns:
# - input.fruits
package filters
@@ -112,7 +112,7 @@ _use_metadata := rego.metadata.chain()`,
rego: `
# METADATA
# scope: package
# custom:
# compile:
# unknowns:
# - input.fruits
package filters
@@ -129,14 +129,14 @@ _use_metadata := rego.metadata.chain()`,
rego: `
# METADATA
# scope: package
# custom:
# compile:
# unknowns:
# - input.fruits
package filters
# METADATA
# scope: document
# custom:
# compile:
# unknowns:
# - input.baskets
include if {
@@ -147,7 +147,7 @@ include if {
# METADATA
# scope: document
# description: if this metadata were picked up, the checks would fail
# custom:
# compile:
# unknowns:
# - input.colour
red_herring if true
@@ -169,14 +169,14 @@ _use_metadata := rego.metadata.chain()`,
rego: `
# METADATA
# scope: package
# custom:
# compile:
# unknowns:
# - input.fruits
package filters
# METADATA
# scope: document
# custom:
# compile:
# unknowns:
# - input.baskets
include if {
@@ -187,7 +187,7 @@ include if {
# METADATA
# scope: document
# description: if this metadata were picked up, the checks would fail
# custom:
# compile:
# unknowns:
# - input.colour
red_herring if true
@@ -281,7 +281,7 @@ red_herring if true
note: "happy path, mapped short unknown",
rego: `package filters
# METADATA
# custom:
# compile:
# unknowns:
# - input.name
include if input.name != "apple"`,
@@ -296,7 +296,7 @@ include if input.name != "apple"`,
note: "happy path, double-mapped short unknown",
rego: `package filters
# METADATA
# custom:
# compile:
# unknowns:
# - input.name
include if input.name != "apple"`,
@@ -701,7 +701,7 @@ other if input.fruits.price > 100
omitUnknowns: true,
rego: `
# METADATA
# custom:
# compile:
# unknowns:
# - inpu.fruits
# - data.whatever
@@ -727,7 +727,7 @@ _use_metadata := rego.metadata.chain()`,
omitUnknowns: true,
rego: `
# METADATA
# custom:
# compile:
# unknowns:
# - inpu.fruits
# - data.whatever
+1 -1
View File
@@ -241,7 +241,7 @@ func TestCompileHandlerHints(t *testing.T) {
typoRego := `package filters
# METADATA
# scope: document
# custom:
# compile:
# unknowns: [input.fruits]
include if input.fruits.name == "apple"
include if input.fruit.cost < input.max
+5 -6
View File
@@ -28,7 +28,7 @@ type FailTracer interface {
Enabled() bool
TraceEvent(topdown.Event)
Config() topdown.TraceConfig
Hints([]*ast.Term) []Hint
Hints([]ast.Ref) []Hint
}
func New() FailTracer {
@@ -54,16 +54,15 @@ func (*failTracer) Config() topdown.TraceConfig {
return topdown.TraceConfig{PlugLocalVars: true}
}
func (b *failTracer) Hints(unknowns []*ast.Term) []Hint {
func (b *failTracer) Hints(unknowns []ast.Ref) []Hint {
var hints []Hint //nolint:prealloc
seenRefs := map[string]struct{}{}
candidates := make([]string, 0, len(unknowns))
for i := range unknowns {
ref, ok := unknowns[i].Value.(ast.Ref)
if !ok || len(ref) < 2 {
for _, ref := range unknowns {
if len(ref) < 2 {
continue
}
candidates = append(candidates, string(unknowns[i].Value.(ast.Ref)[1].Value.(ast.String)))
candidates = append(candidates, string(ref[1].Value.(ast.String)))
}
for _, expr := range b.exprs {
+2 -2
View File
@@ -95,9 +95,9 @@ func TestHints(t *testing.T) {
for i := range tc.evts {
ft.TraceEvent(tc.evts[i])
}
unk := make([]*ast.Term, len(tc.unknowns))
unk := make([]ast.Ref, len(tc.unknowns))
for i := range tc.unknowns {
unk[i] = ast.MustParseTerm(tc.unknowns[i])
unk[i] = ast.MustParseRef(tc.unknowns[i])
}
hints := ft.Hints(unk)
+2 -2
View File
@@ -153,7 +153,7 @@ type Server struct {
cipherSuites *[]uint16
hooks hooks.Hooks
compileUnknownsCache *lru.Cache[string, []*ast.Term]
compileUnknownsCache *lru.Cache[string, []ast.Ref]
compileMaskingRulesCache *lru.Cache[string, ast.Ref]
}
@@ -187,7 +187,7 @@ type Loop func() error
// New returns a new Server.
func New() *Server {
s := Server{}
s.compileUnknownsCache, _ = lru.New[string, []*ast.Term](unknownsCacheSize)
s.compileUnknownsCache, _ = lru.New[string, []ast.Ref](unknownsCacheSize)
s.compileMaskingRulesCache, _ = lru.New[string, ast.Ref](maskingRuleCacheSize)
return &s
}
+1 -1
View File
@@ -1,6 +1,6 @@
# METADATA
# scope: package
# custom:
# compile:
# unknowns:
# - input.tickets
# - input.users