featurebase/cli/writer.go

package cli

import (
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"time"

	featurebase "github.com/featurebasedb/featurebase/v3"
	"github.com/jedib0t/go-pretty/table"
	"github.com/jedib0t/go-pretty/text"
	"github.com/pkg/errors"
)

// writeOptions contains user configuration options which describe how to write
// the query output.
type writeOptions struct {
	border     int
	expanded   bool
	format     string
	location   *time.Location
	timing     bool
	tuplesOnly bool
}

const (
	formatAligned = "aligned"
	formatCSV     = "csv"
)

func defaultWriteOptions() *writeOptions {
	return &writeOptions{
		border:     1,
		expanded:   false,
		format:     formatAligned,
		location:   time.Local,
		timing:     false,
		tuplesOnly: false,
	}
}

// writeOutput writes the query response, taking the format into consideration.
// It sends query output to qOut, non-error informational output (such as query
// timing) to wOut, and errors to wErr.
func writeOutput(r *featurebase.WireQueryResponse, opts *writeOptions, qOut io.Writer, wOut io.Writer, wErr io.Writer) error {
	if r == nil {
		return errors.New("attempt to write out nil response")
	}
	if r.Error != "" {
		if _, err := wErr.Write([]byte("Error: " + r.Error + "\n")); err != nil {
			return errors.Wrapf(err, "writing error: %s", r.Error)
		}
		return writeWarnings(r, wErr)
	}

	switch opts.format {
	case formatAligned:
		if err := writeTable(r, opts, qOut); err != nil {
			return errors.Wrap(err, "writing table")
		}

		// Add some white space after query results.
		qOut.Write([]byte("\n"))

	case formatCSV:
		if err := writeCSV(r, opts, qOut); err != nil {
			return errors.Wrap(err, "writing csv")
		}

	default:
		return errors.Errorf("invalid format: %s", opts.format)
	}

	if err := writeWarnings(r, wErr); err != nil {
		return err
	}

	// Timing.
	if opts.timing {
		if _, err := wOut.Write([]byte(fmt.Sprintf("Execution time: %dμs\n", r.ExecutionTime))); err != nil {
			return errors.Wrapf(err, "writing execution time: %s", r.Error)
		}
	}

	return nil
}

// writeCSV writes the WireQueryResponse to qOut as csv.
func writeCSV(r *featurebase.WireQueryResponse, opts *writeOptions, qOut io.Writer) error {
	w := csv.NewWriter(qOut)

	if opts.expanded {
		// Expanded csv

		// rec is used to write the row as a slice of strings. It is reused to
		// avoid unnecessary memory allocation.
		rec := make([]string, 2)

		for _, row := range r.Data {
			cleanRow(row, opts)
			for i, col := range r.Schema.Fields {
				rec[0] = string(col.Name)
				rec[1] = fmt.Sprintf("%v", row[i])

				// Write the record.
				if err := w.Write(rec); err != nil {
					log.Fatalln("error writing expanded record to csv:", err)
				}
			}
		}

	} else {
		// Normal csv (i.e. NOT expanded)

		// Write the schema.
		if !opts.tuplesOnly {
			header := make([]string, 0, len(r.Schema.Fields))
			for i := range r.Schema.Fields {
				header = append(header, string(r.Schema.Fields[i].Name))
			}
			if err := w.Write(header); err != nil {
				return errors.Wrapf(err, "error writing header to csv")
			}
		}

		// Write the records.

		// rec is used to write the row as a slice of strings. It is reused to
		// avoid unnecessary memory allocation.
		rec := make([]string, len(r.Schema.Fields))

		for _, row := range r.Data {
			cleanRow(row, opts)
			for i := range row {
				rec[i] = fmt.Sprintf("%v", row[i])
			}
			if err := w.Write(rec); err != nil {
				log.Fatalln("error writing record to csv:", err)
			}
		}
	}

	// Write any buffered data to the underlying writer (standard output).
	w.Flush()

	return w.Error()
}

// writeTable writes the WireQueryResponse to qOut in a tabular format.
func writeTable(r *featurebase.WireQueryResponse, opts *writeOptions, qOut io.Writer) error {
	t := table.NewWriter()
	t.SetOutputMirror(qOut)
	switch opts.border {
	case 0:
		t.SetStyle(styleBorder0)
	case 1:
		t.SetStyle(styleBorder1)
	default:
		t.SetStyle(styleBorder2)
		// In expanded mode with a border, we need borders between each record.
		if opts.expanded {
			t.Style().Options.SeparateRows = true
		}
	}

	// Don't uppercase the header values.
	t.Style().Format.Header = text.FormatDefault

	if opts.expanded {
		// Expanded table
		for _, row := range r.Data {
			cleanRow(row, opts)
			colRow := make([]interface{}, 2)
			scolRow := make([]string, 2)
			div := "\n"
			for i, col := range r.Schema.Fields {
				if i == len(r.Schema.Fields)-1 {
					div = ""
				}
				scolRow[0] += fmt.Sprintf("%s%s", col.Name, div)
				scolRow[1] += fmt.Sprintf("%v%s", row[i], div)
			}
			colRow[0] = scolRow[0]
			colRow[1] = scolRow[1]
			t.AppendRow(table.Row(colRow[:]))
		}
	} else {
		// Normal table (i.e. NOT expanded)
		if !opts.tuplesOnly {
			t.AppendHeader(schemaToRow(r.Schema))
		}
		for _, row := range r.Data {
			cleanRow(row, opts)
			t.AppendRow(table.Row(row))
		}
	}
	t.Render()

	return nil
}

// cleanRow loops through all the columns of row and modifies its value based on
// type.
//
// If the value is nil, replace it with a null string; go-pretty doesn't expect
// nil pointers in the data values.
//
// If the value is a time.Time, we want to print it using RFC3339Nano to be
// consistent with everything else.
func cleanRow(row []interface{}, opts *writeOptions) {
	for i := range row {
		switch v := row[i].(type) {
		case nil:
			row[i] = nullValue
		case time.Time:
			row[i] = v.In(opts.location).Format(time.RFC3339Nano)
		}
	}
}

func schemaToRow(schema featurebase.WireQuerySchema) []interface{} {
	ret := make([]interface{}, len(schema.Fields))
	for i, field := range schema.Fields {
		ret[i] = field.Name
	}
	return ret
}

func writeWarnings(r *featurebase.WireQueryResponse, w io.Writer) error {
	if len(r.Warnings) == 0 {
		return nil
	}

	if _, err := w.Write([]byte("\n")); err != nil {
		return errors.Wrapf(err, "writing line feed")
	}
	for _, warning := range r.Warnings {
		if _, err := w.Write([]byte("Warning: " + warning + "\n")); err != nil {
			return errors.Wrapf(err, "writing warning: %s", warning)
		}
	}
	return nil
}

var styleBorder2 table.Style = table.StyleDefault

var styleBorder1 table.Style = table.Style{
	Name:   "StyleBorder1",
	Box:    table.StyleBoxDefault,
	Color:  table.ColorOptionsDefault,
	Format: table.FormatOptionsDefault,
	Options: table.Options{
		DrawBorder:      false,
		SeparateColumns: true,
		SeparateFooter:  true,
		SeparateHeader:  true,
		SeparateRows:    false,
	},
	Title: table.TitleOptionsDefault,
}

var styleBorder0 table.Style = table.Style{
	Name: "StyleBorder0",
	Box: table.BoxStyle{
		BottomLeft:       "+",
		BottomRight:      "+",
		BottomSeparator:  "+",
		Left:             "|",
		LeftSeparator:    "+",
		MiddleHorizontal: "-",
		MiddleSeparator:  " ",
		MiddleVertical:   " ",
		PaddingLeft:      "",
		PaddingRight:     "",
		PageSeparator:    "\n",
		Right:            "|",
		RightSeparator:   "+",
		TopLeft:          "+",
		TopRight:         "+",
		TopSeparator:     "+",
		UnfinishedRow:    " ~",
	},
	Color:  table.ColorOptionsDefault,
	Format: table.FormatOptionsDefault,
	Options: table.Options{
		DrawBorder:      false,
		SeparateColumns: true,
		SeparateFooter:  true,
		SeparateHeader:  true,
		SeparateRows:    false,
	},
	Title: table.TitleOptionsDefault,
}