sq/cli/output/yamlw/recordwriter.go

package yamlw

import (
	"bytes"
	"context"
	"io"
	"strconv"
	"sync"
	"time"

	"github.com/fatih/color"
	goccy "github.com/goccy/go-yaml"
	"github.com/goccy/go-yaml/ast"

	"github.com/neilotoole/sq/cli/output"
	"github.com/neilotoole/sq/libsq/core/errz"
	"github.com/neilotoole/sq/libsq/core/kind"
	"github.com/neilotoole/sq/libsq/core/record"
)

var _ output.NewRecordWriterFunc = NewRecordWriter

// NewRecordWriter returns an output.RecordWriter that writes YAML.
func NewRecordWriter(out io.Writer, pr *output.Printing) output.RecordWriter {
	return &recordWriter{
		out: out,
		pr:  pr,
	}
}

type recordWriter struct {
	mu         sync.Mutex
	out        io.Writer
	pr         *output.Printing
	recMeta    record.Meta
	fieldNames []string
	buf        *bytes.Buffer
	enc        *goccy.Encoder
	clrs       []*color.Color
	keys       []string
	null       string
}

// Open implements output.RecordWriter.
func (w *recordWriter) Open(_ context.Context, recMeta record.Meta) error {
	w.recMeta = recMeta
	w.fieldNames = w.recMeta.MungedNames()
	w.buf = &bytes.Buffer{}
	w.enc = goccy.NewEncoder(io.Discard, decimalMarshaler)
	w.clrs = make([]*color.Color, len(w.recMeta))
	w.keys = make([]string, len(w.recMeta))
	w.null = w.pr.Null.Sprint("null")

	var (
		node ast.Node
		err  error
	)

	// Generate the field keys and colors
	for i := range w.recMeta {
		node, err = w.enc.EncodeToNode(w.recMeta[i].MungedName())
		if err != nil {
			// Shouldn't happen
			return errz.Wrapf(err, "yaml: failed to encode field name: %s", w.recMeta[i].MungedName())
		}
		if node == nil {
			// Also shouldn't happen
			return errz.Errorf("yaml: failed to encode field name: %s: encoded to nil", w.recMeta[i].MungedName())
		}

		w.keys[i] = w.pr.Key.Sprint(node.String())

		// Set up the colors
		switch w.recMeta[i].Kind() {
		case kind.Bytes:
			w.clrs[i] = w.pr.Bytes
		case kind.Null:
			w.clrs[i] = w.pr.Null
		case kind.Text:
			w.clrs[i] = w.pr.String
		case kind.Bool:
			w.clrs[i] = w.pr.Bool
		case kind.Datetime, kind.Date, kind.Time:
			w.clrs[i] = w.pr.Datetime
		case kind.Int, kind.Float, kind.Decimal:
			w.clrs[i] = w.pr.Number
		case kind.Unknown:
			w.clrs[i] = w.pr.Normal
		default:
			w.clrs[i] = w.pr.Normal
		}
	}

	return nil
}

// WriteRecords implements output.RecordWriter.
func (w *recordWriter) WriteRecords(ctx context.Context, recs []record.Record) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	if len(recs) == 0 {
		return nil
	}

	var (
		err  error
		node ast.Node

		buf  = w.buf
		clrs = w.clrs
		keys = w.keys
	)

	for i, rec := range recs {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		buf.WriteString("- ")

		for j := range rec {
			if j > 0 {
				buf.WriteString("\n  ")
			}
			buf.WriteString(keys[j])
			buf.WriteString(": ")
			val := rec[j]

			if val == nil {
				buf.WriteString(w.null)
				continue
			}

			if tm, ok := val.(time.Time); ok {
				if tm.IsZero() {
					buf.WriteString(w.null)
					continue
				}

				var s string
				s, err = w.renderTime(w.recMeta[j], tm)
				if err != nil {
					return err
				}
				buf.WriteString(s)
				continue
			}

			node, err = w.enc.EncodeToNode(val)
			if err != nil {
				return errz.Wrapf(err, "yaml: failed to encode result: [%d].%s",
					i, w.recMeta[j].MungedName())
			}

			if node == nil {
				// Shouldn't happen
				buf.WriteString(w.null)
			} else {
				buf.WriteString(clrs[j].Sprint(node.String()))
			}
		}

		buf.WriteRune('\n')
	}

	return nil
}

// Flush implements output.RecordWriter.
func (w *recordWriter) Flush(context.Context) error {
	w.mu.Lock()
	defer w.mu.Unlock()
	_, err := w.buf.WriteTo(w.out)
	return errz.Err(err)
}

// Close implements output.RecordWriter.
func (w *recordWriter) Close(context.Context) error {
	return w.Flush(context.TODO())
}

// renderTime renders the *time.Time val into a fully-rendered string
// ready for writing out.
func (w *recordWriter) renderTime(fieldMeta *record.FieldMeta, val any) (string, error) {
	if val == nil {
		return w.null, nil
	}

	var (
		pr            = w.pr
		timeFormatter func(time.Time) string
		asNumber      bool
		isNumber      bool
		tm            time.Time
		ok            bool
	)

	if tm, ok = val.(time.Time); !ok {
		return "", errz.Errorf("unexpected value type: expected %T, but got %T", tm, val)
	}

	if tm.IsZero() {
		return w.null, nil
	}

	switch fieldMeta.Kind() { //nolint:exhaustive
	case kind.Datetime:
		timeFormatter = pr.FormatDatetime
		asNumber = pr.FormatDatetimeAsNumber
	case kind.Date:
		timeFormatter = pr.FormatDate
		asNumber = pr.FormatDateAsNumber
	case kind.Time:
		timeFormatter = pr.FormatTime
		asNumber = pr.FormatTimeAsNumber
	default:
		// Shouldn't happen
		return "", errz.Errorf("unexpected data kind: expected a time-like value, but got {%s}: %v",
			fieldMeta.Kind(), val)
	}

	v := timeFormatter(tm)
	if _, err := strconv.ParseInt(v, 10, 64); err == nil {
		isNumber = true
	}

	switch {
	case isNumber && asNumber:
		return w.pr.Datetime.Sprint(v), nil
	case isNumber:
		// If the value is a number, but asNumber is false, we want to render it
		// as a string? Not sure about the desired behavior.

		// We could surround the number in quotes, e.g.
		//
		//  return w.pr.Datetime.Sprintf("%q", v), nil
		//
		// But is that what we want?
		// We'll just leave it naked for now, and await user feedback.
		return w.pr.Datetime.Sprint(v), nil
	default:
		// It's not a number, it's some kind of non-numeric date/time value.
	}

	node, err := w.enc.EncodeToNode(v)
	if err != nil {
		return "", errz.Err(err)
	}

	return w.pr.Datetime.Sprint(node.String()), nil
}