1 files changed, 79 insertions, 125 deletions

diff --git a/internal/mapr/groupset.go b/internal/mapr/groupset.go
index b5c8a48..9d7661a 100644
--- a/internal/mapr/groupset.go
+++ b/internal/mapr/groupset.go
@@ -2,15 +2,9 @@ package mapr
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"io/ioutil"
-	"os"
 	"sort"
 	"strconv"
-	"strings"
-
-	"github.com/mimecast/dtail/internal/io/logger"
 )
 
 // GroupSet represents a map of aggregate sets. The group sets
@@ -22,6 +16,14 @@ type GroupSet struct {
 	sets map[string]*AggregateSet
 }
 
+// Internal helper type
+type result struct {
+	groupKey     string
+	values       []string
+	columnWidths []int
+	orderBy      float64
+}
+
 // NewGroupSet returns a new empty group set.
 func NewGroupSet() *GroupSet {
 	g := GroupSet{}
@@ -56,139 +58,91 @@ func (g *GroupSet) Serialize(ctx context.Context, ch chan<- string) {
 	}
 }
 
-// Result returns a nicely formated result of the query from the group set.
-func (g *GroupSet) Result(query *Query) (string, int, error) {
-	return g.limitedResult(query, query.Limit, "\t", " ", false)
-}
-
-// WriteResult writes the result to an outfile.
-func (g *GroupSet) WriteResult(query *Query) error {
-	if !query.HasOutfile() {
-		return errors.New("No outfile specified")
-	}
-
-	// -1: Don't limit the result, include all data sets
-	result, _, err := g.limitedResult(query, query.Limit, "", ",", true)
-	if err != nil {
-		return err
-	}
-
-	logger.Info("Writing outfile", query.Outfile)
-	tmpOutfile := fmt.Sprintf("%s.tmp", query.Outfile)
-
-	if err := ioutil.WriteFile(tmpOutfile, []byte(result), 0644); err != nil {
-		return err
-	}
-
-	if err := os.Rename(tmpOutfile, query.Outfile); err != nil {
-		os.Remove(tmpOutfile)
-		return err
-	}
-
-	return nil
-}
+// Return a sorted result slice of the query from the group set.
+func (g *GroupSet) result(query *Query, gathercolumnWidths bool) ([]result, []int, error) {
+	var err error
+	var rows []result
 
-// Return a nicely formated result of the query from the group set.
-func (g *GroupSet) limitedResult(query *Query, limit int, lineStarter, fieldSeparator string, addHeader bool) (string, int, error) {
-	type result struct {
-		groupKey  string
-		resultStr string
-		orderBy   float64
-	}
-
-	var resultSlice []result
+	// Helpers for calculating the ASCII table output (output is the terminal and
+	// not a CSV file).
+	columnWidths := make([]int, len(query.Select))
+	var valueStrLen int
 
 	for groupKey, set := range g.sets {
-		var sb strings.Builder
-		r := result{groupKey: groupKey}
+		result := result{groupKey: groupKey}
 
-		lastIndex := len(query.Select) - 1
 		for i, sc := range query.Select {
-			storage := sc.FieldStorage
-			orderByThis := storage == query.OrderBy
-
-			switch sc.Operation {
-			case Count:
-				value := set.FValues[storage]
-				sb.WriteString(fmt.Sprintf("%d", int(value)))
-				if orderByThis {
-					r.orderBy = value
-				}
-			case Len:
-				fallthrough
-			case Sum:
-				fallthrough
-			case Min:
-				fallthrough
-			case Max:
-				value := set.FValues[storage]
-				sb.WriteString(fmt.Sprintf("%f", value))
-				if orderByThis {
-					r.orderBy = value
-				}
-			case Last:
-				value := set.SValues[storage]
-				if orderByThis {
-					f, err := strconv.ParseFloat(value, 64)
-					if err == nil {
-						r.orderBy = f
-					}
-				}
-				sb.WriteString(value)
-			case Avg:
-				value := set.FValues[storage] / float64(set.Samples)
-				sb.WriteString(fmt.Sprintf("%f", value))
-				if orderByThis {
-					r.orderBy = value
-				}
-			default:
-				return "", 0, fmt.Errorf("Unknown aggregation method '%v'", sc.Operation)
+			if valueStrLen, err = g.resultSelect(query, &sc, set, &result); err != nil {
+				return rows, columnWidths, err
+			}
+
+			// Do we want to gather the table withs? This is required to print out a decent
+			// ASCII formated table (table output is the terminal and not a CSV file).
+			if !gathercolumnWidths {
+				continue
+			}
+			if columnWidths[i] < len(sc.FieldStorage) {
+				columnWidths[i] = len(sc.FieldStorage)
 			}
-			if i != lastIndex {
-				sb.WriteString(fieldSeparator)
+			if columnWidths[i] < valueStrLen {
+				columnWidths[i] = valueStrLen
 			}
 		}
+		rows = append(rows, result)
+	}
+
+	g.resultOrderBy(query, rows)
+	return rows, columnWidths, nil
+}
 
-		r.resultStr = sb.String()
-		resultSlice = append(resultSlice, r)
+func (*GroupSet) resultSelect(query *Query, sc *selectCondition, set *AggregateSet,
+	result *result) (int, error) {
+
+	var valueStr string
+	var value float64
+
+	switch sc.Operation {
+	case Count:
+		value = set.FValues[sc.FieldStorage]
+		valueStr = fmt.Sprintf("%d", int(value))
+	case Len:
+		fallthrough
+	case Sum:
+		fallthrough
+	case Min:
+		fallthrough
+	case Max:
+		value = set.FValues[sc.FieldStorage]
+		valueStr = fmt.Sprintf("%f", value)
+	case Last:
+		valueStr = set.SValues[sc.FieldStorage]
+		value, _ = strconv.ParseFloat(valueStr, 64)
+	case Avg:
+		value = set.FValues[sc.FieldStorage] / float64(set.Samples)
+		valueStr = fmt.Sprintf("%f", value)
+	default:
+		return 0, fmt.Errorf("Unknown aggregation method '%v'", sc.Operation)
 	}
 
-	if query.OrderBy != "" {
-		if query.ReverseOrder {
-			sort.SliceStable(resultSlice, func(i, j int) bool {
-				return resultSlice[i].orderBy < resultSlice[j].orderBy
-			})
-		} else {
-			sort.SliceStable(resultSlice, func(i, j int) bool {
-				return resultSlice[i].orderBy > resultSlice[j].orderBy
-			})
-		}
+	if sc.FieldStorage == query.OrderBy {
+		result.orderBy = value
 	}
+	result.values = append(result.values, valueStr)
 
-	var sb strings.Builder
+	return len(valueStr), nil
+}
 
-	// Write header first
-	if addHeader {
-		lastIndex := len(query.Select) - 1
-		sb.WriteString(lineStarter)
-		for i, sc := range query.Select {
-			sb.WriteString(sc.FieldStorage)
-			if i != lastIndex {
-				sb.WriteString(fieldSeparator)
-			}
-		}
-		sb.WriteString("\n")
+func (*GroupSet) resultOrderBy(query *Query, rows []result) {
+	if query.OrderBy == "" {
+		return
 	}
-
-	// And now write the data
-	for i, r := range resultSlice {
-		if i == limit {
-			break
-		}
-		sb.WriteString(lineStarter)
-		sb.WriteString(r.resultStr)
-		sb.WriteString("\n")
+	if query.ReverseOrder {
+		sort.SliceStable(rows, func(i, j int) bool {
+			return rows[i].orderBy < rows[j].orderBy
+		})
+	} else {
+		sort.SliceStable(rows, func(i, j int) bool {
+			return rows[i].orderBy > rows[j].orderBy
+		})
 	}
-
-	return sb.String(), len(resultSlice), nil
 }