Add ExcludeLocationFile JSON marshalling option (#6398)

Repeating the name of the file in each location is often redundant, and large AST trees carry thousands of these attributes. Providing an option to have them removed at least for serialization (as anything more would be a breaking change) seems like a good compromise. For reference, see StyraInc/regal#408 Signed-off-by: Anders Eknert <[email protected]>
open-policy-agent · Nov 10, 2023 · 69a4f3d · 69a4f3d
1 parent d38a4f1
commit 69a4f3d
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 2 deletions.
diff --git a/ast/json/json.go b/ast/json/json.go
@@ -13,6 +13,9 @@ type MarshalOptions struct {
 	IncludeLocation NodeToggle
 	// IncludeLocationText additionally/optionally includes the text of the location
 	IncludeLocationText bool
+	// ExcludeLocationFile additionally/optionally excludes the file of the location
+	// Note that this is inverted (i.e. not "include" as the default needs to remain false)
+	ExcludeLocationFile bool
 }
 
 // NodeToggle is a generic struct to allow the toggling of

diff --git a/ast/location/location.go b/ast/location/location.go
@@ -18,7 +18,7 @@ type Location struct {
 	Col    int    `json:"col"`  // The column in the row.
 	Offset int    `json:"-"`    // The byte offset for the location in the source.
 
-	// JSONOptions specifies options for marshaling and unmarshaling of locations
+	// JSONOptions specifies options for marshaling and unmarshalling of locations
 	JSONOptions astJSON.Options
 }
 
@@ -96,15 +96,32 @@ func (loc *Location) Compare(other *Location) int {
 
 func (loc *Location) MarshalJSON() ([]byte, error) {
 	// structs are used here to preserve the field ordering of the original Location struct
+	if loc.JSONOptions.MarshalOptions.ExcludeLocationFile {
+		data := struct {
+			Row  int    `json:"row"`
+			Col  int    `json:"col"`
+			Text []byte `json:"text,omitempty"`
+		}{
+			Row: loc.Row,
+			Col: loc.Col,
+		}
+
+		if loc.JSONOptions.MarshalOptions.IncludeLocationText {
+			data.Text = loc.Text
+		}
+
+		return json.Marshal(data)
+	}
+
 	data := struct {
 		File string `json:"file"`
 		Row  int    `json:"row"`
 		Col  int    `json:"col"`
 		Text []byte `json:"text,omitempty"`
 	}{
-		File: loc.File,
 		Row:  loc.Row,
 		Col:  loc.Col,
+		File: loc.File,
 	}
 
 	if loc.JSONOptions.MarshalOptions.IncludeLocationText {

diff --git a/ast/location/location_test.go b/ast/location/location_test.go
@@ -116,6 +116,19 @@ func TestLocationMarshal(t *testing.T) {
 			},
 			exp: `{"file":"file","row":1,"col":1,"text":"dGV4dA=="}`,
 		},
+		"excluding file": {
+			loc: &Location{
+				File: "file",
+				Row:  1,
+				Col:  1,
+				JSONOptions: astJSON.Options{
+					MarshalOptions: astJSON.MarshalOptions{
+						ExcludeLocationFile: true,
+					},
+				},
+			},
+			exp: `{"row":1,"col":1}`,
+		},
 	}
 
 	for id, tc := range testCases {

diff --git a/ast/marshal_test.go b/ast/marshal_test.go
@@ -62,6 +62,28 @@ func TestGeneric_MarshalWithLocationJSONOptions(t *testing.T) {
 			}(),
 			ExpectedJSON: `{"location":{"file":"example.rego","row":1,"col":2,"text":"dGhpbmdz"},"type":"string","value":"example"}`,
 		},
+		"location included, location text included, file excluded": {
+			Term: func() *Term {
+				v, _ := InterfaceToValue("example")
+				t := &Term{
+					Value:    v,
+					Location: NewLocation([]byte("things"), "example.rego", 1, 2),
+				}
+				t.setJSONOptions(
+					astJSON.Options{
+						MarshalOptions: astJSON.MarshalOptions{
+							IncludeLocation: astJSON.NodeToggle{
+								Term: true,
+							},
+							IncludeLocationText: true,
+							ExcludeLocationFile: true,
+						},
+					},
+				)
+				return t
+			}(),
+			ExpectedJSON: `{"location":{"row":1,"col":2,"text":"dGhpbmdz"},"type":"string","value":"example"}`,
+		},
 	}
 
 	for name, data := range testCases {

diff --git a/ast/parser_test.go b/ast/parser_test.go
@@ -3959,6 +3959,35 @@ func TestRuleFromBodyJSONOptions(t *testing.T) {
 	}
 }
 
+func TestRuleFromBodyJSONOptionsLocationOptions(t *testing.T) {
+	parserOpts := ParserOptions{ProcessAnnotation: true}
+	parserOpts.JSONOptions = &astJSON.Options{
+		MarshalOptions: astJSON.MarshalOptions{
+			IncludeLocation: astJSON.NodeToggle{
+				Term:           true,
+				Package:        true,
+				Comment:        true,
+				Import:         true,
+				Rule:           true,
+				Head:           true,
+				Expr:           true,
+				SomeDecl:       true,
+				Every:          true,
+				With:           true,
+				Annotations:    true,
+				AnnotationsRef: true,
+			},
+			IncludeLocationText: true,
+			ExcludeLocationFile: true,
+		},
+	}
+
+	module := `package a.b.c
+			foo := "bar"
+			`
+	assertParseModuleJSONOptions(t, `foo := "bar"`, module, parserOpts)
+}
+
 func TestRuleModulePtr(t *testing.T) {
 	mod := `package test