tip/io_writers.go at master · c4pt0r/tip · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package main

import (
	"bufio"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
)

type ResultIOWriter interface {
	Write(rows []RowResult) error
	Flush() error
}

type CSVResultIOWriter struct {
	writer *csv.Writer
}

func NewCSVResultIOWriter(writer io.Writer) *CSVResultIOWriter {
	return &CSVResultIOWriter{
		writer: csv.NewWriter(writer),
	}
}

func (w *CSVResultIOWriter) Write(rows []RowResult) error {
	for _, row := range rows {
		record := make([]string, len(row.colValues))
		for i, val := range row.colValues {
			record[i] = formatCSVValue(val)
		}
		if err := w.writer.Write(record); err != nil {
			return err
		}
	}
	return nil
}

func (w *CSVResultIOWriter) Flush() error {
	w.writer.Flush()
	return w.writer.Error()
}

type PlainResultIOWriter struct {
	writer *bufio.Writer
}

func NewPlainResultIOWriter(writer io.Writer) *PlainResultIOWriter {
	return &PlainResultIOWriter{
		writer: bufio.NewWriter(writer),
	}
}

func (w *PlainResultIOWriter) Write(rows []RowResult) error {
	for _, row := range rows {
		for i, col := range row.colNames {
			val := row.colValues[i]
			_, err := fmt.Fprintf(w.writer, "%s: %s ", col, formatValue(val))
			if err != nil {
				return err
			}
		}
		_, err := w.writer.WriteString("\n")
		if err != nil {
			return err
		}
	}
	return nil
}

func (w *PlainResultIOWriter) Flush() error {
	return w.writer.Flush()
}

type JSONResultIOWriter struct {
	writer *bufio.Writer
	first  bool
}

func NewJSONResultIOWriter(writer io.Writer) *JSONResultIOWriter {
	return &JSONResultIOWriter{
		writer: bufio.NewWriter(writer),
		first:  true,
	}
}

func (w *JSONResultIOWriter) Write(rows []RowResult) error {
	for _, row := range rows {
		if w.first {
			_, err := w.writer.WriteString("[")
			if err != nil {
				return err
			}
			w.first = false
		} else {
			_, err := w.writer.WriteString(",")
			if err != nil {
				return err
			}
		}

		jsonData, err := json.Marshal(row)
		if err != nil {
			return err
		}

		_, err = w.writer.Write(jsonData)
		if err != nil {
			return err
		}
	}
	return nil
}

func (w *JSONResultIOWriter) Flush() error {
	if !w.first {
		_, err := w.writer.WriteString("]")
		if err != nil {
			return err
		}
	}
	return w.writer.Flush()
}