From 077b86c5543c0e9334f710787642714add27e70f Mon Sep 17 00:00:00 2001 From: Miles Maddox Date: Tue, 8 Nov 2022 10:50:43 -0600 Subject: [PATCH] feat: add jsonl output --- lib/query.go | 98 ++++++++++++++++++++++++++--------------------- lib/query_enum.go | 22 ++++++----- 2 files changed, 67 insertions(+), 53 deletions(-) diff --git a/lib/query.go b/lib/query.go index e81fe1c..320ff58 100644 --- a/lib/query.go +++ b/lib/query.go @@ -37,7 +37,7 @@ type Query struct { // Format is an enumeration of available query output formats // ENUM( -// json, csv, table, tsv, xlsx +// json, jsonl, csv, table, tsv, xlsx // ) type Format int @@ -139,8 +139,22 @@ func (q *Query) Execute() (*os.File, error) { func (q *Query) RenderQueryResults(file *os.File) error { var err error - if q.Format == FormatJson.String() { + var outFile *os.File + + if q.OutputFile == "" { + outFile = os.Stdout + } else { + outFile, err = os.Create(q.OutputFile) + if err != nil { + return err + } + defer outFile.Close() + } + + writer := bufio.NewWriter(outFile) + defer writer.Flush() + if q.Format == FormatJson.String() { reader := csvmap.NewReader(file) reader.Columns, err = reader.ReadHeader() if err != nil { @@ -152,15 +166,40 @@ func (q *Query) RenderQueryResults(file *os.File) error { return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err) } - output, _ := json.MarshalIndent(records, "", " ") - - if q.OutputFile == "" { - fmt.Println(string(output)) - } else { - return ioutil.WriteFile(q.OutputFile, output, 0644) + output, err := json.MarshalIndent(records, "", " ") + if err != nil { + return fmt.Errorf("Unable to marshal json %v", err) } + writer.Write(output) return nil + + } + + if q.Format == FormatJsonl.String() { + + reader := csvmap.NewReader(file) + reader.Columns, err = reader.ReadHeader() + if err != nil { + return fmt.Errorf("Unable to read header from %q, %v", file.Name(), err) + } + + for { + record, err := reader.Read() + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err) + } + + output, err := json.Marshal(record) + if err != nil { + return fmt.Errorf("Unable to convert output to jsonl, %v", err) + } + writer.WriteString(string(output) + "\n") + } + } if q.Format == FormatTable.String() { @@ -175,18 +214,7 @@ func (q *Query) RenderQueryResults(file *os.File) error { return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err) } - var f *os.File - - if q.OutputFile == "" { - f = os.Stdout - } else { - f, err = os.Create(q.OutputFile) - if err != nil { - return err - } - } - - table := tablewriter.NewWriter(f) + table := tablewriter.NewWriter(outFile) table.SetHeader(reader.Columns) table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false}) table.SetCenterSeparator("|") @@ -200,33 +228,15 @@ func (q *Query) RenderQueryResults(file *os.File) error { } if q.Format == FormatCsv.String() { - - sb, err := ioutil.ReadFile(file.Name()) + records, err := ioutil.ReadFile(file.Name()) if err != nil { return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err) } - if q.OutputFile == "" { - fmt.Println(string(sb)) - } else { - return ioutil.WriteFile(q.OutputFile, sb, 0644) - } - + writer.Write(records) } if q.Format == FormatTsv.String() { - var w io.Writer - - if q.OutputFile == "" { - w = os.Stdout - } else { - f, err := os.Create(q.OutputFile) - if err != nil { - return fmt.Errorf("Unable to create output file %q, %v", q.OutputFile, err) - } - defer f.Close() - w = bufio.NewWriter(f) - } csvFile, err := os.Open(file.Name()) if err != nil { @@ -235,10 +245,10 @@ func (q *Query) RenderQueryResults(file *os.File) error { defer csvFile.Close() reader := csv.NewReader(csvFile) - writer := csv.NewWriter(w) - defer writer.Flush() + w := csv.NewWriter(writer) + defer w.Flush() - writer.Comma = '\t' + w.Comma = '\t' for { record, err := reader.Read() @@ -249,7 +259,7 @@ func (q *Query) RenderQueryResults(file *os.File) error { return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err) } - if err := writer.Write(record); err != nil { + if err := w.Write(record); err != nil { return fmt.Errorf("error writing record to output %v", err) } } diff --git a/lib/query_enum.go b/lib/query_enum.go index 6450009..6a44fd6 100644 --- a/lib/query_enum.go +++ b/lib/query_enum.go @@ -13,6 +13,8 @@ import ( const ( // FormatJson is a Format of type Json. FormatJson Format = iota + // FormatJsonl is a Format of type Jsonl. + FormatJsonl // FormatCsv is a Format of type Csv. FormatCsv // FormatTable is a Format of type Table. @@ -23,14 +25,15 @@ const ( FormatXlsx ) -const _FormatName = "jsoncsvtabletsvxlsx" +const _FormatName = "jsonjsonlcsvtabletsvxlsx" var _FormatMap = map[Format]string{ FormatJson: _FormatName[0:4], - FormatCsv: _FormatName[4:7], - FormatTable: _FormatName[7:12], - FormatTsv: _FormatName[12:15], - FormatXlsx: _FormatName[15:19], + FormatJsonl: _FormatName[4:9], + FormatCsv: _FormatName[9:12], + FormatTable: _FormatName[12:17], + FormatTsv: _FormatName[17:20], + FormatXlsx: _FormatName[20:24], } // String implements the Stringer interface. @@ -43,10 +46,11 @@ func (x Format) String() string { var _FormatValue = map[string]Format{ _FormatName[0:4]: FormatJson, - _FormatName[4:7]: FormatCsv, - _FormatName[7:12]: FormatTable, - _FormatName[12:15]: FormatTsv, - _FormatName[15:19]: FormatXlsx, + _FormatName[4:9]: FormatJsonl, + _FormatName[9:12]: FormatCsv, + _FormatName[12:17]: FormatTable, + _FormatName[17:20]: FormatTsv, + _FormatName[20:24]: FormatXlsx, } // ParseFormat attempts to convert a string to a Format.