ast: Support inline schema definitions in annotations

Signed-off-by: Torin Sandall <[email protected]>

ast/check.go

@@ -1160,9 +1160,15 @@ func getRuleAnnotation(as *annotationSet, rule *Rule) (result []*SchemaAnnotatio
 
 func processAnnotation(ss *SchemaSet, annot *SchemaAnnotation, env *TypeEnv, rule *Rule) (Ref, types.Type, *Error) {
 
-	schema := ss.Get(annot.Schema)
-	if schema == nil {
-		return nil, nil, NewError(TypeErr, rule.Location, "undefined schema: %v", annot.Schema)
+	var schema interface{}
+
+	if annot.Schema != nil {
+		schema = ss.Get(annot.Schema)
+		if schema == nil {
+			return nil, nil, NewError(TypeErr, rule.Location, "undefined schema: %v", annot.Schema)
+		}
+	} else if annot.Definition != nil {
+		schema = *annot.Definition
 	}
 
 	tpe, err := loadSchema(schema)

ast/check_test.go

@@ -1888,6 +1888,14 @@ package test
 # scope: rule
 # schemas:
 # - input: schema.number
+p { input = 7 }`},
+
+		{note: "inline definition", err: "test.rego:7: rego_type_error: match error", module: `package test
+
+# METADATA
+# scope: rule
+# schemas:
+# - input: {"type": "string"}
 p { input = 7 }`},
 	}
 

ast/compare_test.go

@@ -289,6 +289,41 @@ func TestCompareAnnotations(t *testing.T) {
 # - input: schema.b`,
 			exp: 1,
 		},
+		{
+			note: "definition",
+			a: `
+# METADATA
+# schemas:
+# - input: {"type": "string"}`,
+			b: `
+# METADATA
+# schemas:
+# - input: {"type": "string"}`,
+		},
+		{
+			note: "definition - less than schema",
+			a: `
+# METADATA
+# schemas:
+# - input: {"type": "string"}`,
+			b: `
+# METADATA
+# schemas:
+# - input: schema.a`,
+			exp: -1,
+		},
+		{
+			note: "schema - greater than definition",
+			a: `
+# METADATA
+# schemas:
+# - input: schema.a`,
+			b: `
+# METADATA
+# schemas:
+# - input: {"type": "string"}`,
+			exp: 1,
+		},
 	}
 
 	for _, tc := range tests {

ast/parser.go

@@ -10,7 +10,9 @@ import (
 	"fmt"
 	"io"
 	"math/big"
+	"strings"
 
+	"github.com/pkg/errors"
 	"gopkg.in/yaml.v2"
 
 	"github.com/open-policy-agent/opa/ast/internal/scanner"
@@ -1563,7 +1565,7 @@ type rawAnnotation struct {
 	Schemas []rawSchemaAnnotation `json:"schemas"`
 }
 
-type rawSchemaAnnotation map[string]string
+type rawSchemaAnnotation map[string]interface{}
 
 type metadataParser struct {
 	buf *bytes.Buffer
@@ -1597,21 +1599,36 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 	result.Scope = raw.Scope
 
 	for _, pair := range raw.Schemas {
-		var k, v string
+		var k string
+		var v interface{}
 		for k, v = range pair {
 		}
-		kr, err := ParseRef(k)
+
+		var a SchemaAnnotation
+		var err error
+
+		a.Path, err = ParseRef(k)
 		if err != nil {
 			return nil, fmt.Errorf("invalid document reference")
 		}
-		vr, err := parseSchemaRef(v)
-		if err != nil {
-			return nil, err
+
+		switch v := v.(type) {
+		case string:
+			a.Schema, err = parseSchemaRef(v)
+			if err != nil {
+				return nil, err
+			}
+		case map[interface{}]interface{}:
+			w, err := convertYAMLMapKeyTypes(v, nil)
+			if err != nil {
+				return nil, errors.Wrap(err, "invalid schema definition")
+			}
+			a.Definition = &w
+		default:
+			return nil, fmt.Errorf("invalid schema declaration for path %q", k)
 		}
-		result.Schemas = append(result.Schemas, &SchemaAnnotation{
-			Path:   kr,
-			Schema: vr,
-		})
+
+		result.Schemas = append(result.Schemas, &a)
 	}
 
 	result.Location = b.loc
@@ -1641,3 +1658,32 @@ func parseSchemaRef(s string) (Ref, error) {
 
 	return nil, errInvalidSchemaRef
 }
+
+func convertYAMLMapKeyTypes(x interface{}, path []string) (interface{}, error) {
+	var err error
+	switch x := x.(type) {
+	case map[interface{}]interface{}:
+		result := make(map[string]interface{}, len(x))
+		for k, v := range x {
+			str, ok := k.(string)
+			if !ok {
+				return nil, fmt.Errorf("invalid map key type(s): %v", strings.Join(path, "/"))
+			}
+			result[str], err = convertYAMLMapKeyTypes(v, append(path, str))
+			if err != nil {
+				return nil, err
+			}
+		}
+		return result, nil
+	case []interface{}:
+		for i := range x {
+			x[i], err = convertYAMLMapKeyTypes(x[i], append(path, fmt.Sprintf("%d", i)))
+			if err != nil {
+				return nil, err
+			}
+		}
+		return x, nil
+	default:
+		return x, nil
+	}
+}

ast/parser_test.go

@@ -2672,6 +2672,11 @@ func TestAnnotations(t *testing.T) {
 	schemaNetworks := MustParseRef("schema.networks")
 	schemaPorts := MustParseRef("schema.ports")
 
+	stringSchemaAsMap := map[string]interface{}{
+		"type": "string",
+	}
+	var stringSchema interface{} = stringSchemaAsMap
+
 	tests := []struct {
 		note           string
 		module         string
@@ -3074,6 +3079,24 @@ import data.foo.bar`,
 import data.foo`,
 			expError: "test.rego:2: rego_parse_error: annotation scope 'package' cannot be applied to import statement",
 		},
+		{
+			note: "Inline schema definition",
+			module: `package test
+
+# METADATA
+# schemas:
+# - input: {"type": "string"}
+p { input = "str" }`,
+			expNumComments: 3,
+			expAnnotations: []*Annotations{
+				&Annotations{
+					Schemas: []*SchemaAnnotation{
+						{Path: InputRootRef, Definition: &stringSchema},
+					},
+					Scope: annotationScopeRule,
+				},
+			},
+		},
 	}
 
 	for _, tc := range tests {

ast/policy.go

@@ -154,8 +154,9 @@ type (
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
 	SchemaAnnotation struct {
-		Path   Ref `json:"path"`
-		Schema Ref `json:"schema"`
+		Path       Ref          `json:"path"`
+		Schema     Ref          `json:"schema,omitempty"`
+		Definition *interface{} `json:"definition,omitempty"`
 	}
 
 	// Package represents the namespace of the documents produced
@@ -302,6 +303,14 @@ func (s *SchemaAnnotation) Compare(other *SchemaAnnotation) int {
 		return cmp
 	}
 
+	if s.Definition != nil && other.Definition == nil {
+		return -1
+	} else if s.Definition == nil && other.Definition != nil {
+		return 1
+	} else if s.Definition != nil && other.Definition != nil {
+		return util.Compare(*s.Definition, *other.Definition)
+	}
+
 	return 0
 }