From a8c2ff603d0f8a8494c703487a6e60e2e8aaca9a Mon Sep 17 00:00:00 2001 From: Travis Turner Date: Mon, 13 Mar 2023 17:41:08 -0500 Subject: [PATCH] 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 | 2 +- cli/cli_integration_test.go | 1 + cli/meta.go | 41 ++++++- cli/testdata/meta_pset_format_csv | 47 ++++++++ cli/testdata/meta_timing | 4 +- cli/testdata/setup | 1 + cli/testdata/table | 2 +- cli/writer.go | 173 +++++++++++++++++++++++------- cli/writer_test.go | 7 +- 9 files changed, 231 insertions(+), 47 deletions(-) create mode 100644 cli/testdata/meta_pset_format_csv diff --git a/cli/cli.go b/cli/cli.go index ba4f48bc5..5f902be87 100644 --- a/cli/cli.go +++ b/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") } diff --git a/cli/cli_integration_test.go b/cli/cli_integration_test.go index 0a063e315..1ca7966ec 100644 --- a/cli/cli_integration_test.go +++ b/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", diff --git a/cli/meta.go b/cli/meta.go index 040f5d7c1..a06d7b713 100644 --- a/cli/meta.go +++ b/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) diff --git a/cli/testdata/meta_pset_format_csv b/cli/testdata/meta_pset_format_csv new file mode 100644 index 000000000..18f2417ce --- /dev/null +++ b/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 \ No newline at end of file diff --git a/cli/testdata/meta_timing b/cli/testdata/meta_timing index 9ee2f47e3..736f419cb 100644 --- a/cli/testdata/meta_timing +++ b/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:+-----+-------+-----+ diff --git a/cli/testdata/setup b/cli/testdata/setup index f5c901ae8..4f11aa77b 100644 --- a/cli/testdata/setup +++ b/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 diff --git a/cli/testdata/table b/cli/testdata/table index 0bcccdef7..c2f2e2752 100644 --- a/cli/testdata/table +++ b/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:+-----+-------+-----+ diff --git a/cli/writer.go b/cli/writer.go index ca2f7a2db..2073d0c51 100644 --- a/cli/writer.go +++ b/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{} { diff --git a/cli/writer_test.go b/cli/writer_test.go index 4bfec9627..b89d1d6b6 100644 --- a/cli/writer_test.go +++ b/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())