miller/pkg/output/record_writer_csv.go

package output

import (
	"bufio"
	"fmt"
	"strings"

	csv "github.com/johnkerl/miller/v6/pkg/go-csv"

	"github.com/johnkerl/miller/v6/pkg/cli"
	"github.com/johnkerl/miller/v6/pkg/mlrval"
	"github.com/johnkerl/miller/v6/pkg/types"
)

type RecordWriterCSV struct {
	writerOptions     *cli.TWriterOptions
	csvWriter         *csv.Writer
	needToPrintHeader bool
	firstRecordKeys   []string
	firstRecordNF     int64
	quoteAll          bool // For double-quote around all fields
	// fieldsBuffer is a reusable scratch slice for the stringified field values
	// of the record currently being written. WriteCSVRecordMaybeColorized
	// consumes it synchronously and does not retain it, and the writer processes
	// one record at a time, so a single buffer can be shared across records to
	// avoid a per-record allocation.
	fieldsBuffer []string
}

func NewRecordWriterCSV(writerOptions *cli.TWriterOptions) (*RecordWriterCSV, error) {
	if len(writerOptions.OFS) != 1 {
		return nil, fmt.Errorf("for CSV, OFS can only be a single character")
	}
	if writerOptions.ORS != "\n" && writerOptions.ORS != "\r\n" {
		return nil, fmt.Errorf("for CSV, ORS cannot be altered")
	}
	writer := &RecordWriterCSV{
		writerOptions:     writerOptions,
		csvWriter:         nil, // will be set on first Write() wherein we have the output stream
		needToPrintHeader: !writerOptions.HeaderlessOutput,
		firstRecordKeys:   nil,
		firstRecordNF:     -1,
		quoteAll:          writerOptions.CSVQuoteAll,
	}
	return writer, nil
}

func (writer *RecordWriterCSV) Write(
	outrec *mlrval.Mlrmap,
	_ *types.Context,
	bufferedOutputStream *bufio.Writer,
	outputIsStdout bool,
) error {
	if outrec == nil {
		// End of record stream: nothing special for this output format
		return nil
	}

	if writer.csvWriter == nil {
		writer.csvWriter = csv.NewWriter(bufferedOutputStream)
		writer.csvWriter.Comma = rune(writer.writerOptions.OFS[0]) // xxx temp -- needs length-1 OFS
	}

	if writer.firstRecordKeys == nil {
		writer.firstRecordKeys = outrec.GetKeys()
		writer.firstRecordNF = int64(len(writer.firstRecordKeys))
	}

	if writer.needToPrintHeader {
		fields := make([]string, outrec.FieldCount)
		i := 0
		for pe := outrec.Head; pe != nil; pe = pe.Next {
			fields[i] = pe.Key
			i++
		}
		err := writer.WriteCSVRecordMaybeColorized(fields, bufferedOutputStream, outputIsStdout, true, writer.quoteAll)
		if err != nil {
			return err
		}
		writer.needToPrintHeader = false
	}

	outputNF := outrec.FieldCount
	if outputNF < writer.firstRecordNF {
		outputNF = writer.firstRecordNF
	}

	if int64(cap(writer.fieldsBuffer)) < outputNF {
		writer.fieldsBuffer = make([]string, outputNF)
	}
	fields := writer.fieldsBuffer[:outputNF]
	var i int64 = 0
	for pe := outrec.Head; pe != nil; pe = pe.Next {
		if i < writer.firstRecordNF && pe.Key != writer.firstRecordKeys[i] {
			return fmt.Errorf(
				"CSV schema change: first keys \"%s\"; current keys \"%s\"",
				strings.Join(writer.firstRecordKeys, writer.writerOptions.OFS),
				strings.Join(outrec.GetKeys(), writer.writerOptions.OFS),
			)
		}
		fields[i] = pe.Value.String()
		i++
	}

	for ; i < outputNF; i++ {
		fields[i] = ""
	}

	return writer.WriteCSVRecordMaybeColorized(fields, bufferedOutputStream, outputIsStdout, false, writer.quoteAll)
}