path: root/internal/mapr/groupsetresult.go

package mapr

import (
	"errors"
	"fmt"
	"os"
	"strings"

	"github.com/mimecast/dtail/internal/io/dlog"
	"github.com/mimecast/dtail/internal/io/pool"
	"github.com/mimecast/dtail/internal/protocol"
)

// Result returns a nicely formated result of the query from the group set.
func (g *GroupSet) Result(query *Query, rowsLimit int, renderer ResultRenderer) (string, int, error) {
	rows, columnWidths, err := g.result(query, true)
	if err != nil {
		return "", 0, err
	}
	if query.Limit != -1 {
		rowsLimit = query.Limit
	}
	lastColumn := len(query.Select) - 1

	sb := pool.BuilderBuffer.Get().(*strings.Builder)
	defer pool.RecycleBuilderBuffer(sb)

	if renderer == nil {
		renderer = PlainResultRenderer()
	}

	g.resultWriteFormattedHeader(query, renderer, sb, lastColumn, rowsLimit, columnWidths)
	g.resultWriteFormattedHeaderRowSeparator(query, renderer, sb, lastColumn, columnWidths)
	g.resultWriteFormattedData(query, renderer, sb, lastColumn, rowsLimit, columnWidths, rows)

	return sb.String(), len(rows), nil
}

// Write a nicely formatted header for the result data.
func (g *GroupSet) resultWriteFormattedHeader(query *Query, renderer ResultRenderer, sb *strings.Builder,
	lastColumn, rowsLimit int, columnWidths []int) {

	for i, sc := range query.Select {
		format := fmt.Sprintf(" %%%ds ", columnWidths[i])
		str := fmt.Sprintf(format, sc.FieldStorage)

		g.resultWriteFormattedHeaderEntry(query, renderer, sb, sc, str)
		if i == lastColumn {
			continue
		}
		g.resultWriteFormattedHeaderEntrySeparator(renderer, sb)

	}
	sb.WriteString("\n")
}

func (g *GroupSet) resultWriteFormattedHeaderEntry(query *Query, renderer ResultRenderer, sb *strings.Builder,
	sc selectCondition, str string) {

	isGroupKey := false
	for _, groupBy := range query.GroupBy {
		if sc.FieldStorage == groupBy {
			isGroupKey = true
			break
		}
	}
	renderer.WriteHeaderEntry(sb, str, sc.FieldStorage == query.OrderBy, isGroupKey)
}

func (g *GroupSet) resultWriteFormattedHeaderEntrySeparator(renderer ResultRenderer, sb *strings.Builder) {
	renderer.WriteHeaderDelimiter(sb, protocol.FieldDelimiter)
}

// This writes a nicely formatted line separating the header and the data.
func (g *GroupSet) resultWriteFormattedHeaderRowSeparator(query *Query, renderer ResultRenderer, sb *strings.Builder,
	lastColumn int, columnWidths []int) {

	for i := 0; i < len(query.Select); i++ {
		str := fmt.Sprintf("-%s-", strings.Repeat("-", columnWidths[i]))
		renderer.WriteHeaderDelimiter(sb, str)
		if i == lastColumn {
			continue
		}
		renderer.WriteHeaderDelimiter(sb, protocol.FieldDelimiter)
	}
	sb.WriteString("\n")
}

// Write the result data nicely formatted.
func (g *GroupSet) resultWriteFormattedData(query *Query, renderer ResultRenderer, sb *strings.Builder,
	lastColumn, rowsLimit int, columnWidths []int, rows []result) {

	for i, r := range rows {
		if i == rowsLimit {
			break
		}
		for j, value := range r.values {
			g.resultWriteFormattedDataEntry(renderer, sb, columnWidths, j, value)
			if j == lastColumn {
				continue
			}
			renderer.WriteDataDelimiter(sb, protocol.FieldDelimiter)
		}
		sb.WriteString("\n")
	}
}

func (g *GroupSet) resultWriteFormattedDataEntry(renderer ResultRenderer, sb *strings.Builder,
	columnWidths []int, j int, value string) {

	format := fmt.Sprintf(" %%%ds ", columnWidths[j])
	str := fmt.Sprintf(format, value)
	renderer.WriteDataEntry(sb, str)
}

func (*GroupSet) writeQueryFile(query *Query) error {
	queryFile := fmt.Sprintf("%s.query", query.Outfile.FilePath)
	tmpQueryFile := fmt.Sprintf("%s.tmp", queryFile)
	dlog.Common.Debug("Writing query file", queryFile)

	fd, err := os.OpenFile(tmpQueryFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	if err != nil {
		return err
	}
	defer fd.Close()

	if _, err := fd.WriteString(query.RawQuery); err != nil {
		return err
	}
	return os.Rename(tmpQueryFile, queryFile)
}

// WriteResult writes the result to an CSV outfile.
func (g *GroupSet) WriteResult(query *Query, finalResult bool) error {
	if !query.HasOutfile() {
		return errors.New("No outfile specified")
	}
	if err := g.writeQueryFile(query); err != nil {
		return err
	}
	rows, _, err := g.result(query, false)
	if err != nil {
		return err
	}

	// By default, also write the CSV header.
	writeHeader := true

	// In append mode, only write CSV header when file doesn't exist yet or is empty.
	if query.Outfile.AppendMode {
		if info, err := os.Stat(query.Outfile.FilePath); err == nil && info.Size() > 0 {
			writeHeader = false
		}
	}

	fd, err := g.getOutfileFD(query)
	if err != nil {
		return err
	}
	defer fd.Close()

	return g.resultWriteUnformatted(query, rows, fd, writeHeader, finalResult)
}

func (g *GroupSet) getOutfileFD(query *Query) (*os.File, error) {
	if !query.Outfile.AppendMode {
		dlog.Common.Info("Writing to outfile", query.Outfile.FilePath)
		tmpOutfile := fmt.Sprintf("%s.tmp", query.Outfile.FilePath)
		return os.OpenFile(tmpOutfile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	}

	dlog.Common.Info("Appending to outfile", query.Outfile.FilePath)
	return os.OpenFile(query.Outfile.FilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
}

func (g *GroupSet) resultWriteUnformatted(query *Query, rows []result, fd *os.File, writeHeader, finalResult bool) error {
	lastColumn := len(query.Select) - 1

	if writeHeader {
		if err := g.resultWriteUnformattedHeader(query, fd, lastColumn); err != nil {
			return err
		}
	}

	// And now write the data
	for i, r := range rows {
		if i == query.Limit {
			break
		}
		for j, value := range r.values {
			if _, err := fd.WriteString(value); err != nil {
				return err
			}
			if j == lastColumn {
				continue
			}
			if _, err := fd.WriteString(protocol.CSVDelimiter); err != nil {
				return err
			}
		}
		if _, err := fd.WriteString("\n"); err != nil {
			return err
		}
	}

	// Always rename .tmp to .csv after writing (not just on final result)
	// This ensures the .csv file is updated at every interval
	if !query.Outfile.AppendMode {
		tmpOutfile := fmt.Sprintf("%s.tmp", query.Outfile.FilePath)
		dlog.Common.Debug("Renaming outfile", tmpOutfile, "to", query.Outfile.FilePath)
		if err := os.Rename(tmpOutfile, query.Outfile.FilePath); err != nil {
			dlog.Common.Error("Failed to rename outfile", tmpOutfile, "error", err)
			os.Remove(tmpOutfile)
			return err
		}
		dlog.Common.Info("Successfully renamed outfile to", query.Outfile.FilePath)
	}

	return nil
}

func (g *GroupSet) resultWriteUnformattedHeader(query *Query, fd *os.File, lastColumn int) (err error) {
	for i, sc := range query.Select {
		if _, err = fd.WriteString(sc.FieldStorage); err != nil {
			return
		}
		if i == lastColumn {
			continue
		}
		if _, err = fd.WriteString(protocol.CSVDelimiter); err != nil {
			return
		}
	}
	_, err = fd.WriteString("\n")
	return
}