Add CSV support to fbsql (#2318)

* Pre csv cleanup

* Implement the CSV writer

This adds the `format` sub-command to `\pset` in order to choose between
formats `aligned` and `csv`.

cli/cli.go

@@ -393,7 +393,7 @@ func (cmd *Command) executeAndWriteQuery(qry query) error {
 		}
 		return errors.Wrap(err, "making query")
 	}
-	if err := writeTable(queryResponse, cmd.writeOptions, cmd.output, cmd.stdout, cmd.stderr); err != nil {
+	if err := writeOutput(queryResponse, cmd.writeOptions, cmd.output, cmd.stdout, cmd.stderr); err != nil {
 		return errors.Wrap(err, "writing out response")
 	}
 

cli/cli_integration_test.go

@@ -69,6 +69,7 @@ func TestCLIIntegration(t *testing.T) {
 			"meta_file",
 			"meta_pset_border",
 			"meta_pset_expanded",
+			"meta_pset_format_csv",
 			"meta_pset_tuples_only",
 			"meta_include",
 			"meta_output",

cli/meta.go

@@ -293,6 +293,40 @@ func (m *metaFile) execute(cmd *Command) (responseAction, error) {
 	return actionNone, errors.Wrap(err, "adding part file")
 }
 
+// ////////////////////////////////////////////////////////////////////////////
+// format (sub-command of pset)
+// ////////////////////////////////////////////////////////////////////////////
+type metaFormat struct {
+	args []string
+}
+
+func newMetaFormat(args []string) *metaFormat {
+	return &metaFormat{
+		args: args,
+	}
+}
+
+func (m *metaFormat) execute(cmd *Command) (responseAction, error) {
+	switch len(m.args) {
+	case 0:
+		// pass
+	case 1:
+		switch m.args[0] {
+		case formatAligned:
+			cmd.writeOptions.format = formatAligned
+		case formatCSV:
+			cmd.writeOptions.format = formatCSV
+		default:
+			return actionNone, errors.Errorf("\\pset: allowed formats are aligned, csv")
+		}
+	default:
+		return actionNone, errors.Errorf("meta command 'format' takes zero or one argument")
+	}
+
+	cmd.Printf("Output format is %s.\n", cmd.writeOptions.format)
+	return actionNone, nil
+}
+
 // ////////////////////////////////////////////////////////////////////////////
 // help (?)
 // ////////////////////////////////////////////////////////////////////////////
@@ -335,7 +369,7 @@ Informational
 
 Formatting
   \pset [NAME [VALUE]]   set table output option
-                         (border|expanded|location|tuples_only)
+                         (border|expanded|format|location|tuples_only)
   \t [on|off]            show only rows
   \x [on|off]            toggle expanded output
 
@@ -634,6 +668,7 @@ func (m *metaPSet) print(cmd *Command) {
 
 	fmt := `border      %d
 expanded    %s
+format      %s
 location    %s
 tuples_only %s
 `
@@ -641,6 +676,7 @@ tuples_only %s
 	cmd.Printf(fmt,
 		cmd.writeOptions.border,
 		onOff(cmd.writeOptions.expanded),
+		cmd.writeOptions.format,
 		cmd.writeOptions.location,
 		onOff(cmd.writeOptions.tuplesOnly),
 	)
@@ -660,6 +696,9 @@ func (m *metaPSet) execute(cmd *Command) (responseAction, error) {
 		case "expanded", "x":
 			sub := newMetaExpanded(m.args[1:])
 			return sub.execute(cmd)
+		case "format":
+			sub := newMetaFormat(m.args[1:])
+			return sub.execute(cmd)
 		case "location":
 			sub := newMetaLocation(m.args[1:])
 			return sub.execute(cmd)

cli/testdata/meta_pset_format_csv

@@ -0,0 +1,47 @@
+SEND:\pset format csv
+EXPECT:Output format is csv.
+
+SEND:SELECT * FROM users;
+EXPECT:_id,name,age
+EXPECT:1,Anne,38
+EXPECT:2,Bill,23
+EXPECT:3,Cindy,64
+
+// Exclude headers.
+SEND:\t on
+EXPECT:Tuples only is on.
+
+SEND:SELECT * FROM users;
+EXPECT:1,Anne,38
+EXPECT:2,Bill,23
+EXPECT:3,Cindy,64
+
+// Reset headers.
+SEND:\t off
+EXPECT:Tuples only is off.
+
+// Set expanded to on.
+SEND:\x on
+EXPECT:Expanded display is on.
+
+SEND:SELECT * FROM users;
+EXPECT:_id,1
+EXPECT:name,Anne
+EXPECT:age,38
+EXPECT:_id,2
+EXPECT:name,Bill
+EXPECT:age,23
+EXPECT:_id,3
+EXPECT:name,Cindy
+EXPECT:age,64
+
+// Set expanded back to off.
+SEND:\x off
+EXPECT:Expanded display is off.
+
+// Set format back to aligned as we started.
+SEND:\pset format aligned
+EXPECT:Output format is aligned.
+
+SEND:\pset format invalid
+EXPECT:executing meta command: \pset: allowed formats are aligned, csv

cli/testdata/meta_timing

@@ -19,7 +19,7 @@ SEND:\timing on extra
 EXPECT:executing meta command: meta command 'timing' takes zero or one argument
 
 
-SEND:SELECT * from users;
+SEND:SELECT * FROM users;
 EXPECT:+-----+-------+-----+
 EXPECT:| _id | name  | age |
 EXPECT:+-----+-------+-----+
@@ -35,7 +35,7 @@ SEND:\timing off
 EXPECT:Timing is off.
 
 // Ensure we don't get timing.
-SEND:SELECT * from users;
+SEND:SELECT * FROM users;
 EXPECT:+-----+-------+-----+
 EXPECT:| _id | name  | age |
 EXPECT:+-----+-------+-----+

cli/testdata/setup

@@ -46,5 +46,6 @@ EXPECT:Border style is 2.
 SEND:\pset
 EXPECT:border      2
 EXPECT:expanded    off
+EXPECT:format      aligned
 EXPECT:location    UTC
 EXPECT:tuples_only off

cli/testdata/table

@@ -49,7 +49,7 @@ EXPECT:
 // running this creates the fb_views sytem table which has a description.
 // And it's annoying to mask out all of the description fields because we
 // don't know which row fb_views will fall into.
-SEND:SELECT * from users;
+SEND:SELECT * FROM users;
 EXPECT:+-----+-------+-----+
 EXPECT:| _id | name  | age |
 EXPECT:+-----+-------+-----+

cli/writer.go

@@ -1,8 +1,10 @@
 package cli
 
 import (
+	"encoding/csv"
 	"fmt"
 	"io"
+	"log"
 	"time"
 
 	featurebase "github.com/featurebasedb/featurebase/v3"
@@ -16,25 +18,32 @@ import (
 type writeOptions struct {
 	border     int
 	expanded   bool
+	format     string
+	location   *time.Location
 	timing     bool
 	tuplesOnly bool
-	location   *time.Location
 }
 
+const (
+	formatAligned = "aligned"
+	formatCSV     = "csv"
+)
+
 func defaultWriteOptions() *writeOptions {
 	return &writeOptions{
 		border:     1,
 		expanded:   false,
+		format:     formatAligned,
+		location:   time.Local,
 		timing:     false,
 		tuplesOnly: false,
-		location:   time.Local,
 	}
 }
 
-// writeTable writes the query response, taking the format into consideration.
+// 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 writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.Writer, wOut io.Writer, wErr io.Writer) error {
+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")
 	}
@@ -45,9 +54,104 @@ func writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.
 		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 format.border {
+	switch opts.border {
 	case 0:
 		t.SetStyle(styleBorder0)
 	case 1:
@@ -55,7 +159,7 @@ func writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.
 	default:
 		t.SetStyle(styleBorder2)
 		// In expanded mode with a border, we need borders between each record.
-		if format.expanded {
+		if opts.expanded {
 			t.Style().Options.SeparateRows = true
 		}
 	}
@@ -63,9 +167,10 @@ func writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.
 	// Don't uppercase the header values.
 	t.Style().Format.Header = text.FormatDefault
 
-	if format.expanded {
+	if opts.expanded {
 		// Expanded table
 		for _, row := range r.Data {
+			cleanRow(row, opts)
 			colRow := make([]interface{}, 2)
 			scolRow := make([]string, 2)
 			div := "\n"
@@ -74,11 +179,7 @@ func writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.
 					div = ""
 				}
 				scolRow[0] += fmt.Sprintf("%s%s", col.Name, div)
-				if row[i] == nil {
-					scolRow[1] += fmt.Sprintf("%s%s", nullValue, div)
-				} else {
-					scolRow[1] += fmt.Sprintf("%v%s", row[i], div)
-				}
+				scolRow[1] += fmt.Sprintf("%v%s", row[i], div)
 			}
 			colRow[0] = scolRow[0]
 			colRow[1] = scolRow[1]
@@ -86,46 +187,36 @@ func writeTable(r *featurebase.WireQueryResponse, format *writeOptions, qOut io.
 		}
 	} else {
 		// Normal table (i.e. NOT expanded)
-		if !format.tuplesOnly {
+		if !opts.tuplesOnly {
 			t.AppendHeader(schemaToRow(r.Schema))
 		}
 		for _, row := range r.Data {
-			// Loop through all the colums of each row and modify any 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.
-			for i := range row {
-				switch v := row[i].(type) {
-				case nil:
-					row[i] = nullValue
-				case time.Time:
-					row[i] = v.In(format.location).Format(time.RFC3339Nano)
-				}
-			}
+			cleanRow(row, opts)
 			t.AppendRow(table.Row(row))
 		}
 	}
 	t.Render()
 
-	if err := writeWarnings(r, wErr); err != nil {
-		return err
-	}
+	return nil
+}
 
-	// Add some white space after query results.
-	qOut.Write([]byte("\n"))
-
-	// Timing.
-	if format.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)
+// 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)
 		}
 	}
-
-	return nil
 }
 
 func schemaToRow(schema featurebase.WireQuerySchema) []interface{} {

cli/writer_test.go

@@ -58,6 +58,7 @@ func TestWriter(t *testing.T) {
 				format: &writeOptions{
 					border:     1,
 					expanded:   false,
+					format:     formatAligned,
 					timing:     true,
 					tuplesOnly: false,
 				},
@@ -77,6 +78,7 @@ func TestWriter(t *testing.T) {
 				format: &writeOptions{
 					border:     2,
 					expanded:   false,
+					format:     formatAligned,
 					timing:     false,
 					tuplesOnly: false,
 				},
@@ -98,6 +100,7 @@ func TestWriter(t *testing.T) {
 				format: &writeOptions{
 					border:     0,
 					expanded:   false,
+					format:     formatAligned,
 					timing:     false,
 					tuplesOnly: false,
 				},
@@ -117,6 +120,7 @@ func TestWriter(t *testing.T) {
 				format: &writeOptions{
 					border:     1,
 					expanded:   false,
+					format:     formatAligned,
 					timing:     false,
 					tuplesOnly: true,
 				},
@@ -134,6 +138,7 @@ func TestWriter(t *testing.T) {
 				format: &writeOptions{
 					border:     2,
 					expanded:   true,
+					format:     formatAligned,
 					timing:     false,
 					tuplesOnly: false,
 				},
@@ -164,7 +169,7 @@ func TestWriter(t *testing.T) {
 				wOut := bytes.NewBuffer(make([]byte, 0, 100000))
 				wErr := bytes.NewBuffer(make([]byte, 0, 100000))
 
-				assert.NoError(t, writeTable(wqr, test.format, qOut, wOut, wErr))
+				assert.NoError(t, writeOutput(wqr, test.format, qOut, wOut, wErr))
 
 				assert.Equal(t, test.expQOut, qOut.String())
 				assert.Equal(t, test.expOut, wOut.String())