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 }