Add pretty-printing option for query results

cmd/fleetctl/query.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -12,15 +11,11 @@ import (
 	"github.com/urfave/cli"
 )
 
-type resultOutput struct {
-	HostIdentifier string              `json:"host"`
-	Rows           []map[string]string `json:"rows"`
-}
 
 func queryCommand() cli.Command {
 	var (
 		flHosts, flLabels, flQuery, flQueryName string
-		flDebug, flQuiet, flExit                bool
+		flDebug, flQuiet, flExit, flPretty      bool
 		flTimeout                               time.Duration
 	)
 	return cli.Command{
@@ -76,6 +71,12 @@ func queryCommand() cli.Command {
 				Destination: &flDebug,
 				Usage:       "Whether or not to enable debug logging",
 			},
+			cli.BoolFlag{
+				Name:        "pretty",
+				EnvVar:      "PRETTY",
+				Destination: &flPretty,
+				Usage:       "Enable pretty-printing",
+			},
 			cli.DurationFlag{
 				Name:        "timeout",
 				EnvVar:      "TIMEOUT",
@@ -109,6 +110,13 @@ func queryCommand() cli.Command {
 				return fmt.Errorf("Query must be specified with --query or --query-name")
 			}
 
+			var output outputWriter
+			if flPretty {
+				output = newPrettyWriter()
+			} else {
+				output = newJsonWriter()
+			}
+
 			hosts := strings.Split(flHosts, ",")
 			labels := strings.Split(flLabels, ",")
 
@@ -142,11 +150,12 @@ func queryCommand() cli.Command {
 				select {
 				// Print a result
 				case hostResult := <-res.Results():
-					out := resultOutput{hostResult.Host.HostName, hostResult.Rows}
 					s.Stop()
-					if err := json.NewEncoder(os.Stdout).Encode(out); err != nil {
-						fmt.Fprintf(os.Stderr, "Error writing output: %s\n", err)
+
+					if err := output.WriteResult(hostResult); err != nil {
+						fmt.Fprintf(os.Stderr, "Error writing result: %s\n", err)
 					}
+
 					s.Start()
 
 				// Print an error

cmd/fleetctl/query_output.go

@@ -0,0 +1,83 @@
+package main
+
+import (
+	"sort"
+	"encoding/json"
+	"os"
+
+	"github.com/olekukonko/tablewriter"
+	"github.com/gosuri/uilive"
+	"github.com/kolide/fleet/server/kolide"
+)
+
+type outputWriter interface{
+	WriteResult(res kolide.DistributedQueryResult) error
+}
+
+type resultOutput struct {
+	HostIdentifier string              `json:"host"`
+	Rows           []map[string]string `json:"rows"`
+}
+
+type jsonWriter struct {}
+
+func newJsonWriter() *jsonWriter {
+	return &jsonWriter{}
+}
+
+func (w *jsonWriter) WriteResult(res kolide.DistributedQueryResult) error {
+	out := resultOutput{res.Host.HostName, res.Rows}
+	return json.NewEncoder(os.Stdout).Encode(out)
+}
+
+type prettyWriter struct {
+	results []kolide.DistributedQueryResult
+	columns map[string]bool
+	writer *uilive.Writer
+}
+
+func newPrettyWriter() *prettyWriter{
+	return &prettyWriter{
+		columns: make(map[string]bool),
+		writer: uilive.New(),
+	}
+}
+
+func (w *prettyWriter) WriteResult(res kolide.DistributedQueryResult) error {
+	w.results = append(w.results, res)
+
+	// Recompute columns
+	for _, row := range res.Rows {
+		delete(row, "host_hostname")
+		for col := range row {
+			w.columns[col] = true
+		}
+	}
+
+	columns := []string{}
+	for col := range w.columns {
+		columns = append(columns, col)
+	}
+	sort.Strings(columns)
+
+	table := tablewriter.NewWriter(w.writer.Newline())
+	table.SetRowLine(true)
+	table.SetHeader(append([]string{"hostname"}, columns...))
+
+	// Extract columns from the results in the appropriate order
+	for _, res := range w.results {
+		for _, row := range res.Rows {
+			cols := []string{res.Host.HostName}
+			for _, col := range columns {
+				cols = append(cols, row[col])
+			}
+			table.Append(cols)
+		}
+	}
+	table.Render()
+
+	// Actually write the output
+	w.writer.Flush()
+
+	return nil
+}

go.mod

@@ -23,6 +23,7 @@ require (
 	github.com/google/uuid v0.0.0-20161128191214-064e2069ce9c // indirect
 	github.com/gorilla/mux v1.6.2
 	github.com/gorilla/websocket v1.4.2
+	github.com/gosuri/uilive v0.0.4
 	github.com/hashicorp/golang-lru v0.5.1 // indirect
 	github.com/hashicorp/hcl v0.0.0-20180404174102-ef8a98b0bbce // indirect
 	github.com/igm/sockjs-go v3.0.0+incompatible

go.sum

@@ -81,6 +81,8 @@ github.com/gorilla/mux v1.6.2 h1:Pgr17XVTNXAk3q/r4CpKzC5xBM/qW1uVLV+IhRZpIIk=
 github.com/gorilla/mux v1.6.2/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs=
 github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/gosuri/uilive v0.0.4 h1:hUEBpQDj8D8jXgtCdBu7sWsy5sbW/5GhuO8KBwJ2jyY=
+github.com/gosuri/uilive v0.0.4/go.mod h1:V/epo5LjjlDE5RJUcqx8dbw+zc93y5Ya3yg8tfZ74VI=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1 h1:0hERBMJE1eitiLkihrMvRVBYAkpHzc/J3QdDN+dAcgU=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=