You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
1.9 KiB
91 lines
1.9 KiB
|
2 months ago
|
package mock
|
||
|
|
|
||
|
|
import (
|
||
|
|
"encoding/csv"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"sort"
|
||
|
|
"strings"
|
||
|
|
)
|
||
|
|
|
||
|
|
func logExportFormat(value string) (string, error) {
|
||
|
|
value = strings.ToLower(strings.TrimSpace(value))
|
||
|
|
if value == "" {
|
||
|
|
value = "jsonl"
|
||
|
|
}
|
||
|
|
if value != "jsonl" && value != "json" && value != "csv" {
|
||
|
|
return "", fmt.Errorf("format must be csv, json, or jsonl")
|
||
|
|
}
|
||
|
|
return value, nil
|
||
|
|
}
|
||
|
|
func writeLogExport(w io.Writer, format string, rows []any) error {
|
||
|
|
format, err := logExportFormat(format)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if format == "json" {
|
||
|
|
return json.NewEncoder(w).Encode(rows)
|
||
|
|
}
|
||
|
|
if format == "jsonl" {
|
||
|
|
e := json.NewEncoder(w)
|
||
|
|
for _, row := range rows {
|
||
|
|
if err := e.Encode(row); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
objects := make([]map[string]json.RawMessage, 0, len(rows))
|
||
|
|
keys := map[string]bool{}
|
||
|
|
for _, row := range rows {
|
||
|
|
b, err := json.Marshal(row)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
var object map[string]json.RawMessage
|
||
|
|
if err = json.Unmarshal(b, &object); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
objects = append(objects, object)
|
||
|
|
for key := range object {
|
||
|
|
keys[key] = true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
columns := make([]string, 0, len(keys))
|
||
|
|
for key := range keys {
|
||
|
|
columns = append(columns, key)
|
||
|
|
}
|
||
|
|
sort.Strings(columns)
|
||
|
|
if _, err := io.WriteString(w, "\ufeff"); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
writer := csv.NewWriter(w)
|
||
|
|
if err := writer.Write(columns); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
for _, object := range objects {
|
||
|
|
values := make([]string, len(columns))
|
||
|
|
for i, key := range columns {
|
||
|
|
raw := object[key]
|
||
|
|
if len(raw) == 0 {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
var value string
|
||
|
|
if json.Unmarshal(raw, &value) != nil {
|
||
|
|
value = string(raw)
|
||
|
|
}
|
||
|
|
trimmed := strings.TrimLeft(value, " \t\r\n")
|
||
|
|
if len(trimmed) > 0 && strings.ContainsRune("=+-@", rune(trimmed[0])) {
|
||
|
|
value = "'" + value
|
||
|
|
}
|
||
|
|
values[i] = value
|
||
|
|
}
|
||
|
|
if err := writer.Write(values); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
writer.Flush()
|
||
|
|
return writer.Error()
|
||
|
|
}
|