sq/cli/output/yamlw/recordwriter.go

254 lines
5.5 KiB
Go
Raw Normal View History

2023-05-11 05:03:45 +03:00
package yamlw
import (
"bytes"
"context"
2023-05-11 05:03:45 +03:00
"io"
"strconv"
"sync"
"time"
"github.com/fatih/color"
goccy "github.com/goccy/go-yaml"
2023-05-11 05:03:45 +03:00
"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"
2023-05-11 05:03:45 +03:00
)
var _ output.NewRecordWriterFunc = NewRecordWriter
2023-05-11 05:03:45 +03:00
// 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
2023-05-11 05:03:45 +03:00
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 {
2023-05-11 05:03:45 +03:00
w.recMeta = recMeta
w.fieldNames = w.recMeta.MungedNames()
2023-05-11 05:03:45 +03:00
w.buf = &bytes.Buffer{}
w.enc = goccy.NewEncoder(io.Discard, decimalMarshaler)
2023-05-11 05:03:45 +03:00
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())
2023-05-11 05:03:45 +03:00
if err != nil {
// Shouldn't happen
return errz.Wrapf(err, "yaml: failed to encode field name: %s", w.recMeta[i].MungedName())
2023-05-11 05:03:45 +03:00
}
if node == nil {
// Also shouldn't happen
return errz.Errorf("yaml: failed to encode field name: %s: encoded to nil", w.recMeta[i].MungedName())
2023-05-11 05:03:45 +03:00
}
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 {
2023-05-11 05:03:45 +03:00
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:
}
2023-05-11 05:03:45 +03:00
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() {
2023-05-11 05:03:45 +03:00
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())
2023-05-11 05:03:45 +03:00
}
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 {
2023-05-11 05:03:45 +03:00
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())
2023-05-11 05:03:45 +03:00
}
// 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) {
2023-05-11 05:03:45 +03:00
if val == nil {
return w.null, nil
}
var (
pr = w.pr
timeFormatter func(time.Time) string
asNumber bool
isNumber bool
tm time.Time
2023-05-11 05:03:45 +03:00
ok bool
)
if tm, ok = val.(time.Time); !ok {
2023-05-11 05:03:45 +03:00
return "", errz.Errorf("unexpected value type: expected %T, but got %T", tm, val)
}
if tm.IsZero() {
return w.null, nil
}
2023-05-11 05:03:45 +03:00
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)
2023-05-11 05:03:45 +03:00
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
}