master
go 202 lines 4.5 KB
Raw
1 // SPDX-License-Identifier: GPL-3.0-or-later
2
3 package logs
4
5 import (
6 "bytes"
7 "encoding/csv"
8 "errors"
9 "fmt"
10 "io"
11 "strconv"
12 "strings"
13 )
14
15 type CSVCheckFieldFunc func(string) (string, int, bool)
16
17 func (f CSVCheckFieldFunc) IsZero() bool { return true }
18
19 type (
20 CSVConfig struct {
21 FieldsPerRecord int `yaml:"fields_per_record,omitempty" json:"fields_per_record"`
22 Delimiter string `yaml:"delimiter,omitempty" json:"delimiter"`
23 TrimLeadingSpace bool `yaml:"trim_leading_space,omitempty" json:"trim_leading_space"`
24 Format string `yaml:"format,omitempty" json:"format"`
25 CheckField CSVCheckFieldFunc `yaml:"-" json:"-"`
26 }
27
28 CSVParser struct {
29 Config CSVConfig
30 reader *csv.Reader
31 format *csvFormat
32 }
33
34 csvFormat struct {
35 raw string
36 maxIndex int
37 fields []csvField
38 }
39
40 csvField struct {
41 name string
42 idx int
43 }
44 )
45
46 func NewCSVParser(config CSVConfig, in io.Reader) (*CSVParser, error) {
47 if config.Format == "" {
48 return nil, errors.New("empty csv format")
49 }
50
51 format, err := newCSVFormat(config)
52 if err != nil {
53 return nil, fmt.Errorf("bad csv format '%s': %v", config.Format, err)
54 }
55
56 p := &CSVParser{
57 Config: config,
58 reader: newCSVReader(in, config),
59 format: format,
60 }
61 return p, nil
62 }
63
64 func (p *CSVParser) ReadLine(line LogLine) error {
65 record, err := p.reader.Read()
66 if err != nil {
67 return handleCSVReaderError(err)
68 }
69 return p.format.parse(record, line)
70 }
71
72 func (p *CSVParser) Parse(row []byte, line LogLine) error {
73 r := newCSVReader(bytes.NewBuffer(row), p.Config)
74 record, err := r.Read()
75 if err != nil {
76 return handleCSVReaderError(err)
77 }
78 return p.format.parse(record, line)
79 }
80
81 func (p CSVParser) Info() string {
82 return fmt.Sprintf("csv: %s", p.format.raw)
83 }
84
85 func (f *csvFormat) parse(record []string, line LogLine) error {
86 if len(record) <= f.maxIndex {
87 return &ParseError{msg: "csv parse: unmatched line"}
88 }
89
90 for _, v := range f.fields {
91 if err := line.Assign(v.name, record[v.idx]); err != nil {
92 return &ParseError{msg: fmt.Sprintf("csv parse: %v", err), err: err}
93 }
94 }
95 return nil
96 }
97
98 func newCSVReader(in io.Reader, config CSVConfig) *csv.Reader {
99 r := csv.NewReader(in)
100 if config.Delimiter != "" {
101 if d, err := parseCSVDelimiter(config.Delimiter); err == nil {
102 r.Comma = d
103 }
104 }
105 r.TrimLeadingSpace = config.TrimLeadingSpace
106 r.FieldsPerRecord = config.FieldsPerRecord
107 r.ReuseRecord = true
108 return r
109 }
110
111 func newCSVFormat(config CSVConfig) (*csvFormat, error) {
112 r := csv.NewReader(strings.NewReader(config.Format))
113 if config.Delimiter != "" {
114 if d, err := parseCSVDelimiter(config.Delimiter); err == nil {
115 r.Comma = d
116 }
117 }
118 r.TrimLeadingSpace = config.TrimLeadingSpace
119
120 record, err := r.Read()
121 if err != nil {
122 return nil, err
123 }
124
125 fields, err := createCSVFields(record, config.CheckField)
126 if err != nil {
127 return nil, err
128 }
129
130 if len(fields) == 0 {
131 return nil, errors.New("zero fields")
132 }
133
134 format := &csvFormat{
135 raw: config.Format,
136 maxIndex: fields[len(fields)-1].idx,
137 fields: fields,
138 }
139 return format, nil
140 }
141
142 func createCSVFields(format []string, check func(string) (string, int, bool)) ([]csvField, error) {
143 if check == nil {
144 check = checkCSVFormatField
145 }
146 var fields []csvField
147 var offset int
148 seen := make(map[string]bool)
149
150 for i, name := range format {
151 name = strings.Trim(name, `"`)
152
153 name, addOffset, valid := check(name)
154 offset += addOffset
155 if !valid {
156 continue
157 }
158 if seen[name] {
159 return nil, fmt.Errorf("duplicate field: %s", name)
160 }
161 seen[name] = true
162
163 idx := i + offset
164 fields = append(fields, csvField{name, idx})
165 }
166 return fields, nil
167 }
168
169 func handleCSVReaderError(err error) error {
170 if isCSVParseError(err) {
171 return &ParseError{msg: fmt.Sprintf("csv parse: %v", err), err: err}
172 }
173 return err
174 }
175
176 func isCSVParseError(err error) bool {
177 return errors.Is(err, csv.ErrBareQuote) || errors.Is(err, csv.ErrFieldCount) || errors.Is(err, csv.ErrQuote)
178 }
179
180 func checkCSVFormatField(name string) (newName string, offset int, valid bool) {
181 if len(name) < 2 || !strings.HasPrefix(name, "$") {
182 return "", 0, false
183 }
184 return name, 0, true
185 }
186
187 func parseCSVDelimiter(s string) (rune, error) {
188 if isNumber(s) {
189 d, err := strconv.ParseInt(s, 10, 32)
190 if err != nil {
191 return 0, fmt.Errorf("invalid CSV delimiter: %v", err)
192 }
193 if d < 0 {
194 return 0, errors.New("invalid CSV delimiter: must be a non-negative integer")
195 }
196 return rune(d), nil
197 }
198 if len(s) != 1 {
199 return 0, errors.New("invalid CSV delimiter: must be a single character")
200 }
201 return rune(s[0]), nil
202 }