feat: implement csv-noheader and tsv-noheader output formats

mergestat · Mar 8, 2022 · 1fe8e56 · 1fe8e56
1 parent 1ea9d18
commit 1fe8e56
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 13 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -28,7 +28,7 @@ var logger = zerolog.Nop()                            // By default use a NOOP l
 
 func init() {
 	// local (root command only) flags
-	rootCmd.Flags().StringVarP(&format, "format", "f", "table", "specify the output format. Options are 'csv' 'tsv' 'table' 'single' 'ndjson' and 'json'")
+	rootCmd.Flags().StringVarP(&format, "format", "f", "table", "specify the output format. Options are 'csv' 'csv-noheader' 'tsv' 'tsv-noheader' 'table' 'single' 'ndjson' and 'json'")
 	rootCmd.Flags().StringVarP(&presetQuery, "preset", "p", "", "used to pick a preset query")
 	rootCmd.PersistentFlags().StringVarP(&dbPath, "db", "d", "", "specify a db file on disk to mount when executing queries")
 	rootCmd.PersistentFlags().StringVarP(&repo, "repo", "r", ".", "specify a path to a default repo on disk. This will be used if no repo is supplied as an argument to a git table")

diff --git a/pkg/display/display.go b/pkg/display/display.go
@@ -18,12 +18,22 @@ func WriteTo(rows *sql.Rows, w io.Writer, format string, interactive bool) error
 			return err
 		}
 	case "csv":
-		err := csvDisplay(rows, ',', w)
+		err := csvDisplay(rows, ',', false, w)
+		if err != nil {
+			return err
+		}
+	case "csv-noheader":
+		err := csvDisplay(rows, ',', true, w)
 		if err != nil {
 			return err
 		}
 	case "tsv":
-		err := csvDisplay(rows, '\t', w)
+		err := csvDisplay(rows, '\t', false, w)
+		if err != nil {
+			return err
+		}
+	case "tsv-noheader":
+		err := csvDisplay(rows, '\t', true, w)
 		if err != nil {
 			return err
 		}
@@ -37,7 +47,6 @@ func WriteTo(rows *sql.Rows, w io.Writer, format string, interactive bool) error
 		if err != nil {
 			return err
 		}
-	//TODO: switch between table and csv dependent on num columns(suggested num for table 5<=
 	default:
 		err := tableDisplay(rows, w, interactive)
 		if err != nil {
@@ -78,18 +87,21 @@ func single(rows *sql.Rows, write io.Writer) error {
 	return nil
 }
 
-func csvDisplay(rows *sql.Rows, commaChar rune, writer io.Writer) error {
+func csvDisplay(rows *sql.Rows, commaChar rune, noHeader bool, writer io.Writer) error {
 	columns, err := rows.Columns()
 	if err != nil {
 		return err
 	}
 	w := csv.NewWriter(writer)
 	w.Comma = commaChar
 
-	err = w.Write(columns)
-	if err != nil {
-		return err
+	if !noHeader {
+		err = w.Write(columns)
+		if err != nil {
+			return err
+		}
 	}
+
 	pointers := make([]interface{}, len(columns))
 	container := make([]sql.NullString, len(columns))
 

diff --git a/pkg/display/display_test.go b/pkg/display/display_test.go
@@ -42,19 +42,24 @@ func TestDisplayTable(t *testing.T) {
 }
 
 func TestDisplayCSV(t *testing.T) {
-	db, mock, _ := sqlmock.New()
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	mockRows := sqlmock.NewRows([]string{"id", "name", "value"}).
 		AddRow("1", "name 1", "value 1").
 		AddRow("2", "name 2", "value 2").
 		AddRow("3", "name 3", "value 3")
 
 	mock.ExpectQuery("select").WillReturnRows(mockRows)
-
-	rows, _ := db.Query("select")
+	rows, err := db.Query("select")
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	var b bytes.Buffer
-	err := WriteTo(rows, &b, "csv", false)
+	err = WriteTo(rows, &b, "csv", false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -75,10 +80,50 @@ func TestDisplayCSV(t *testing.T) {
 	if len(records[0]) != 3 {
 		t.Fatalf("expected 3 columns of output, got: %d", len(records[0]))
 	}
-
 	// TODO perhaps test the actual content of the lines?
 }
 
+func TestDisplayCSVNoHeader(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mockRows := sqlmock.NewRows([]string{"id", "name", "value"}).
+		AddRow("1", "name 1", "value 1").
+		AddRow("2", "name 2", "value 2").
+		AddRow("3", "name 3", "value 3")
+
+	mock.ExpectQuery("select").WillReturnRows(mockRows)
+	rows, err := db.Query("select")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var b bytes.Buffer
+	err = WriteTo(rows, &b, "csv-noheader", false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := csv.NewReader(strings.NewReader(b.String()))
+
+	records, err := r.ReadAll()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	lineCount := len(records)
+
+	if lineCount != 3 {
+		t.Fatalf("expected 3 lines of output, got: %d", lineCount)
+	}
+
+	if len(records[0]) != 3 {
+		t.Fatalf("expected 3 columns of output, got: %d", len(records[0]))
+	}
+}
+
 func TestDisplayTSV(t *testing.T) {
 	db, mock, _ := sqlmock.New()
 
@@ -118,6 +163,45 @@ func TestDisplayTSV(t *testing.T) {
 	// TODO perhaps test the actual content of the lines?
 }
 
+func TestDisplayTSVNoHeader(t *testing.T) {
+	db, mock, _ := sqlmock.New()
+
+	mockRows := sqlmock.NewRows([]string{"id", "name", "value"}).
+		AddRow("1", "name 1", "value 1").
+		AddRow("2", "name 2", "value 2").
+		AddRow("3", "name 3", "value 3")
+
+	mock.ExpectQuery("select").WillReturnRows(mockRows)
+
+	rows, _ := db.Query("select")
+
+	var b bytes.Buffer
+	err := WriteTo(rows, &b, "tsv-noheader", false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := csv.NewReader(strings.NewReader(b.String()))
+	r.Comma = '\t'
+
+	records, err := r.ReadAll()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	lineCount := len(records)
+
+	if lineCount != 3 {
+		t.Fatalf("expected 3 lines of output, got: %d", lineCount)
+	}
+
+	if len(records[0]) != 3 {
+		t.Fatalf("expected 3 columns of output, got: %d", len(records[0]))
+	}
+
+	// TODO perhaps test the actual content of the lines?
+}
+
 func TestDisplayJSON(t *testing.T) {
 	db, mock, _ := sqlmock.New()