mapreduce tables are in colors now too

3 files changed, 45 insertions, 37 deletions

diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index 63ceaac..51f33c1 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"bytes"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -17,7 +18,7 @@ type baseHandler struct {
 	server       string
 	shellStarted bool
 	commands     chan string
-	receiveBuf   []byte
+	receiveBuf   bytes.Buffer
 	status       int
 }
 
@@ -56,14 +57,23 @@ func (h *baseHandler) SendMessage(command string) error {
 // Read data from the dtail server via Writer interface.
 func (h *baseHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		if b == protocol.MessageDelimiter || b == '\n' {
-			if len(h.receiveBuf) == 0 {
+		switch b {
+		/*
+			// TODO: Next DTail version make it so that '\n' gets ignored. For now
+			// leave it for compatibility with older DTail server + ability to display
+			// the protocol mismatch warn message.
+			case '\n' {
+				continue
+		*/
+		case '\n', protocol.MessageDelimiter:
+			message := h.receiveBuf.String()
+			if len(message) == 0 {
 				continue
 			}
-			message := string(h.receiveBuf)
 			h.handleMessageType(message)
-		} else {
-			h.receiveBuf = append(h.receiveBuf, b)
+			h.receiveBuf.Reset()
+		default:
+			h.receiveBuf.WriteByte(b)
 		}
 	}
 
@@ -78,25 +88,22 @@ func (h *baseHandler) Read(p []byte) (n int, err error) {
 	case <-h.Done():
 		return 0, io.EOF
 	}
-
 	return
 }
 
 // Handle various message types.
 func (h *baseHandler) handleMessageType(message string) {
-	if len(h.receiveBuf) == 0 {
+	if len(message) == 0 {
 		return
 	}
 
 	// Hidden server commands starti with a dot "."
-	if h.receiveBuf[0] == '.' {
+	if message[0] == '.' {
 		h.handleHiddenMessage(message)
-		h.receiveBuf = h.receiveBuf[:0]
 		return
 	}
 
 	logger.Raw(message)
-	h.receiveBuf = h.receiveBuf[:0]
 }
 
 // Handle messages received from server which are not meant to be displayed
diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
index 213748c..eca0348 100644
--- a/internal/clients/handlers/healthhandler.go
+++ b/internal/clients/handlers/healthhandler.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"time"
@@ -13,7 +14,7 @@ import (
 type HealthHandler struct {
 	done *internal.Done
 	// Buffer of incoming data from server.
-	receiveBuf []byte
+	receiveBuf bytes.Buffer
 	// To send commands to the server.
 	commands chan string
 	// To receive messages from the server.
@@ -72,10 +73,10 @@ func (h *HealthHandler) SendMessage(command string) error {
 // Server writes byte stream to client.
 func (h *HealthHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		h.receiveBuf = append(h.receiveBuf, b)
-		if b == protocol.MessageDelimiter { // '\n' {
-			h.receive <- string(h.receiveBuf)
-			h.receiveBuf = h.receiveBuf[:0]
+		h.receiveBuf.WriteByte(b)
+		if b == protocol.MessageDelimiter {
+			h.receive <- h.receiveBuf.String()
+			h.receiveBuf.Reset()
 		}
 	}
 
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
index afad507..65b1454 100644
--- a/internal/clients/handlers/maprhandler.go
+++ b/internal/clients/handlers/maprhandler.go
@@ -15,7 +15,6 @@ type MaprHandler struct {
 	baseHandler
 	aggregate *client.Aggregate
 	query     *mapr.Query
-	count     uint64
 }
 
 // NewMaprHandler returns a new mapreduce client handler.
@@ -36,19 +35,20 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr
 // Read data from the dtail server via Writer interface.
 func (h *MaprHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		h.baseHandler.receiveBuf = append(h.baseHandler.receiveBuf, b)
-		if b == protocol.MessageDelimiter { // '\n' {
-			if len(h.baseHandler.receiveBuf) == 0 {
-				continue
+		switch b {
+		case '\n':
+			continue
+		case protocol.MessageDelimiter:
+			message := h.baseHandler.receiveBuf.String()
+			logger.Debug(message)
+			if message[0] == 'A' {
+				h.handleAggregateMessage(message)
+			} else {
+				h.baseHandler.handleMessageType(message)
 			}
-			message := string(h.baseHandler.receiveBuf)
-
-			if h.baseHandler.receiveBuf[0] == 'A' {
-				h.handleAggregateMessage(strings.TrimSpace(message))
-				h.baseHandler.receiveBuf = h.baseHandler.receiveBuf[:0]
-				continue
-			}
-			h.baseHandler.handleMessageType(message)
+			h.baseHandler.receiveBuf.Reset()
+		default:
+			h.baseHandler.receiveBuf.WriteByte(b)
 		}
 	}
 
@@ -58,12 +58,12 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) {
 // Handle a message received from server including mapr aggregation
 // related data.
 func (h *MaprHandler) handleAggregateMessage(message string) {
-	h.count++
-	parts := strings.Split(message, protocol.AggregateDelimiter)
-
-	// Index 0 contains 'AGGREGATE', 1 contains server host.
-	// Aggregation data begins from index 2.
-	logger.Debug("Received aggregate data", h.server, h.count, parts)
-	h.aggregate.Aggregate(parts[2:])
-	logger.Debug("Aggregated aggregate data", h.server, h.count)
+	parts := strings.SplitN(message, protocol.FieldDelimiter, 3)
+	if len(parts) != 3 {
+		logger.Error("Unable to aggregate data", h.server, message, parts, len(parts), "expected 3 parts")
+		return
+	}
+	if err := h.aggregate.Aggregate(parts[2]); err != nil {
+		logger.Error("Unable to aggregate data", h.server, message, err)
+	}
 }