From 0687fe3dbd504e0c2784f2f771f4135a7e1fea18 Mon Sep 17 00:00:00 2001 From: Torin Sandall Date: Thu, 19 May 2016 12:17:15 -0700 Subject: [PATCH] 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 | 1 + runtime/repl.go | 99 +++++++++++++++++++++++++++++--------------- runtime/repl_test.go | 52 ++++++++++++++++++++++- runtime/runtime.go | 20 +++++---- 4 files changed, 130 insertions(+), 42 deletions(-) diff --git a/cmd/run.go b/cmd/run.go index e9571ad7c0..982aabb129 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -73,6 +73,7 @@ In addition, API calls to delete policies will remove the definition file. runCommand.Flags().StringVarP(¶ms.HistoryPath, "history", "H", historyPath(), "set path of history file") runCommand.Flags().StringVarP(¶ms.PolicyDir, "policy-dir", "p", "", "set directory to store policy definitions") runCommand.Flags().StringVarP(¶ms.Addr, "addr", "a", defaultAddr, "set listening address of the server") + runCommand.Flags().StringVarP(¶ms.OutputFormat, "format", "f", "pretty", "set shell output format, i.e, pretty, json") wrapFlags(runCommand.Flags()) flag.Parse() diff --git a/runtime/repl.go b/runtime/repl.go index 16d4e2d1b8..eba25dd3d5 100644 --- a/runtime/repl.go +++ b/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 }{ {"", "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) diff --git a/runtime/repl_test.go b/runtime/repl_test.go index 670850627e..dd80616686 100644 --- a/runtime/repl_test.go +++ b/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 } diff --git a/runtime/runtime.go b/runtime/runtime.go index dd4f151d5b..0b62738bec 100644 --- a/runtime/runtime.go +++ b/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() }