Add knob to control output format in REPL

Add support for JSON output format in addition to pretty format. Pretty
format is nice but sometimes data sets/queries are not well suited for it.

cmd/run.go

@@ -73,6 +73,7 @@ In addition, API calls to delete policies will remove the definition file.
 	runCommand.Flags().StringVarP(&params.HistoryPath, "history", "H", historyPath(), "set path of history file")
 	runCommand.Flags().StringVarP(&params.PolicyDir, "policy-dir", "p", "", "set directory to store policy definitions")
 	runCommand.Flags().StringVarP(&params.Addr, "addr", "a", defaultAddr, "set listening address of the server")
+	runCommand.Flags().StringVarP(&params.OutputFormat, "format", "f", "pretty", "set shell output format, i.e, pretty, json")
 
 	wrapFlags(runCommand.Flags())
 	flag.Parse()

runtime/repl.go

@@ -27,6 +27,7 @@ const (
 // Repl represeents an instance of the interactive shell.
 type Repl struct {
 	Output       io.Writer
+	OutputFormat string
 	Trace        bool
 	Runtime      *Runtime
 	HistoryPath  string
@@ -37,9 +38,10 @@ type Repl struct {
 }
 
 // NewRepl creates a new Repl.
-func NewRepl(rt *Runtime, historyPath string, output io.Writer) *Repl {
+func NewRepl(rt *Runtime, historyPath string, output io.Writer, outputFormat string) *Repl {
 	return &Repl{
 		Output:       output,
+		OutputFormat: outputFormat,
 		Trace:        false,
 		Runtime:      rt,
 		HistoryPath:  historyPath,
@@ -90,6 +92,10 @@ func (r *Repl) OneShot(line string) bool {
 		switch strings.TrimSpace(strings.ToLower(line)) {
 		case "dump":
 			return r.cmdDump()
+		case "json":
+			return r.cmdFormat("json")
+		case "pretty":
+			return r.cmdFormat("pretty")
 		case "trace":
 			return r.cmdTrace()
 		case "?":
@@ -119,6 +125,11 @@ func (r *Repl) cmdDump() bool {
 }
 
 func (r *Repl) cmdExit() bool {
+	return true
+}
+
+func (r *Repl) cmdFormat(s string) bool {
+	r.OutputFormat = s
 	return false
 }
 
@@ -129,11 +140,13 @@ func (r *Repl) cmdHelp() bool {
 		note string
 	}{
 		{"<stmt>", "evaluate the statement"},
+		{"json", "set output format to JSON"},
+		{"pretty", "set output format to pretty"},
 		{"dump", "dump the raw storage content"},
 		{"trace", "toggle stdout tracing"},
-		{"ctrl+l", "clear the screen"},
 		{"help", "print this message (or ?)"},
 		{"exit", "exit back to shell (or ctrl+c, ctrl+d, quit)"},
+		{"ctrl+l", "clear the screen"},
 	}
 
 	maxLength := 0
@@ -334,16 +347,63 @@ func (r *Repl) evalBody(body ast.Body) bool {
 	return false
 }
 
+func (r *Repl) evalRule(rule *ast.Rule) bool {
+
+	path := []interface{}{string(rule.Name)}
+
+	if err := r.Runtime.DataStore.Patch(storage.AddOp, path, []*ast.Rule{rule}); err != nil {
+		fmt.Fprintln(r.Output, "error:", err)
+		return true
+	}
+
+	fmt.Fprintln(r.Output, "defined")
+	return false
+}
+
+func (r *Repl) getPrompt() string {
+	if len(r.Buffer) > 0 {
+		return r.BufferPrompt
+	}
+	return r.InitPrompt
+}
+
+func (r *Repl) loadHistory(prompt *liner.State) {
+	if f, err := os.Open(r.HistoryPath); err == nil {
+		prompt.ReadHistory(f)
+		f.Close()
+	}
+}
+
 func (r *Repl) printResults(body ast.Body, results []map[string]interface{}) {
+
+	switch r.OutputFormat {
+	case "json":
+		r.printJSON(results)
+	default:
+		r.printPretty(body, results)
+	}
+
+}
+
+func (r *Repl) printJSON(results []map[string]interface{}) {
+	buf, err := json.MarshalIndent(results, "", "  ")
+	if err != nil {
+		fmt.Fprintln(r.Output, err)
+		return
+	}
+	fmt.Fprintln(r.Output, string(buf))
+}
+
+func (r *Repl) printPretty(body ast.Body, results []map[string]interface{}) {
 	table := tablewriter.NewWriter(r.Output)
-	r.printHeader(table, body)
+	r.printPrettyHeader(table, body)
 	for _, row := range results {
-		r.printRow(table, row)
+		r.printPrettyRow(table, row)
 	}
 	table.Render()
 }
 
-func (r *Repl) printHeader(table *tablewriter.Table, body ast.Body) {
+func (r *Repl) printPrettyHeader(table *tablewriter.Table, body ast.Body) {
 
 	// Build set of fields for the output. The fields are the variables from inside the body.
 	// If the variable appears multiple times, we only want a single field so store them in a
@@ -372,7 +432,7 @@ func (r *Repl) printHeader(table *tablewriter.Table, body ast.Body) {
 	table.SetHeader(keys)
 }
 
-func (r *Repl) printRow(table *tablewriter.Table, row map[string]interface{}) {
+func (r *Repl) printPrettyRow(table *tablewriter.Table, row map[string]interface{}) {
 
 	// Arrange fields in same order as header.
 	keys := []string{}
@@ -396,33 +456,6 @@ func (r *Repl) printRow(table *tablewriter.Table, row map[string]interface{}) {
 	table.Append(buf)
 }
 
-func (r *Repl) evalRule(rule *ast.Rule) bool {
-
-	path := []interface{}{string(rule.Name)}
-
-	if err := r.Runtime.DataStore.Patch(storage.AddOp, path, []*ast.Rule{rule}); err != nil {
-		fmt.Fprintln(r.Output, "error:", err)
-		return true
-	}
-
-	fmt.Fprintln(r.Output, "defined")
-	return false
-}
-
-func (r *Repl) getPrompt() string {
-	if len(r.Buffer) > 0 {
-		return r.BufferPrompt
-	}
-	return r.InitPrompt
-}
-
-func (r *Repl) loadHistory(prompt *liner.State) {
-	if f, err := os.Open(r.HistoryPath); err == nil {
-		prompt.ReadHistory(f)
-		f.Close()
-	}
-}
-
 func (r *Repl) saveHistory(prompt *liner.State) {
 	if f, err := os.Create(r.HistoryPath); err == nil {
 		prompt.WriteHistory(f)

runtime/repl_test.go

@@ -75,6 +75,56 @@ func TestOneShotBufferedRule(t *testing.T) {
 	expectOutput(t, buffer.String(), "defined\n")
 }
 
+func TestOneShotJSON(t *testing.T) {
+	store := newTestStorage()
+	var buffer bytes.Buffer
+	repl := newRepl(store, &buffer)
+	repl.OutputFormat = "json"
+	repl.OneShot("data.a[i] = x")
+	var expected interface{}
+	input := `
+	[
+		{
+			"i": 0,
+			"x": {
+			"b": {
+				"c": [
+				true,
+				2,
+				false
+				]
+			}
+			}
+		},
+		{
+			"i": 1,
+			"x": {
+			"b": {
+				"c": [
+				false,
+				true,
+				1
+				]
+			}
+			}
+		}
+	]
+	`
+	if err := json.Unmarshal([]byte(input), &expected); err != nil {
+		panic(err)
+	}
+
+	var result interface{}
+
+	if err := json.Unmarshal(buffer.Bytes(), &result); err != nil {
+		t.Errorf("Unexpected output format: %v", err)
+		return
+	}
+	if !reflect.DeepEqual(expected, result) {
+		t.Errorf("Expected %v but got: %v", expected, result)
+	}
+}
+
 func TestBuildHeader(t *testing.T) {
 	expr := ast.MustParseStatement(`[{"a": x, "b": data.a.b[y]}] = [{"a": 1, "b": 2}]`).(ast.Body)[0]
 	terms := expr.Terms.([]*ast.Term)
@@ -96,7 +146,7 @@ func expectOutput(t *testing.T, output string, expected string) {
 
 func newRepl(store *storage.DataStore, buffer *bytes.Buffer) *Repl {
 	runtime := &Runtime{DataStore: store}
-	repl := NewRepl(runtime, "", buffer)
+	repl := NewRepl(runtime, "", buffer, "")
 	return repl
 }
 

runtime/runtime.go

@@ -23,22 +23,26 @@ type Params struct {
 	// Addr is the listening address that the OPA server will bind to.
 	Addr string
 
-	// Server flag controls whether the OPA instance will start a server.
-	// By default, the OPA instance acts as an interactive shell.
-	Server bool
+	// HistoryPath is the filename to store the interactive shell user
+	// input history.
+	HistoryPath string
+
+	// Output format controls how the REPL will print query results.
+	// Default: "pretty".
+	OutputFormat string
 
 	// Paths contains filenames of base documents and policy modules to
 	// load on startup.
 	Paths []string
 
-	// HistoryPath is the filename to store the interactive shell user
-	// input history.
-	HistoryPath string
-
 	// PolicyDir is the filename of the directory to persist policy
 	// definitions in. Policy definitions stored in this directory
 	// are automatically loaded on startup.
 	PolicyDir string
+
+	// Server flag controls whether the OPA instance will start a server.
+	// By default, the OPA instance acts as an interactive shell.
+	Server bool
 }
 
 // Runtime represents a single OPA instance.
@@ -113,7 +117,7 @@ func (rt *Runtime) startServer(params *Params) {
 }
 
 func (rt *Runtime) startRepl(params *Params) {
-	repl := NewRepl(rt, params.HistoryPath, os.Stdout)
+	repl := NewRepl(rt, params.HistoryPath, os.Stdout, params.OutputFormat)
 	repl.Loop()
 }