From db1abbe65a72d47e7e5bb5280c920f6d170071d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Toma=C5=BE=20Jerman?= Date: Sun, 19 Feb 2023 15:16:55 +0100 Subject: [PATCH] Prepare JSON & CSV decoders for large imports --- server/pkg/envoyx/csv/decoder.go | 149 +++++++++++++++++++++++++ server/pkg/envoyx/csv/decoder_test.go | 72 ++++++++++++ server/pkg/envoyx/json/decoder.go | 145 ++++++++++++++++++++++++ server/pkg/envoyx/json/decoder_test.go | 71 ++++++++++++ 4 files changed, 437 insertions(+) create mode 100644 server/pkg/envoyx/csv/decoder.go create mode 100644 server/pkg/envoyx/csv/decoder_test.go create mode 100644 server/pkg/envoyx/json/decoder.go create mode 100644 server/pkg/envoyx/json/decoder_test.go diff --git a/server/pkg/envoyx/csv/decoder.go b/server/pkg/envoyx/csv/decoder.go new file mode 100644 index 000000000..9cf372b7b --- /dev/null +++ b/server/pkg/envoyx/csv/decoder.go @@ -0,0 +1,149 @@ +package csv + +import ( + "encoding/csv" + "io" + "io/ioutil" + "os" + "strings" + + "github.com/gabriel-vasile/mimetype" +) + +type ( + decoder struct { + ident string + + src *os.File + + reader *csv.Reader + skipHead bool + + header []string + row map[string]string + count uint64 + } +) + +func CanDecodeFile(f io.Reader) bool { + m, err := mimetype.DetectReader(f) + if err != nil { + return false + } + + return CanDecodeExt(m.Extension()) +} + +func CanDecodeMime(m string) bool { + return m == "text/csv" +} + +func CanDecodeExt(ext string) bool { + pt := strings.Split(ext, ".") + return strings.TrimSpace(pt[len(pt)-1]) == "csv" +} + +// Decoder inits a new csv decoder from the given reader +// +// @todo hold small files in mem to avoid needles disc access +func Decoder(r io.Reader, ident string) (out *decoder, err error) { + out = &decoder{ + ident: ident, + } + + out.src, err = ioutil.TempFile(os.TempDir(), "*.csv") + if err != nil { + return + } + + r, err = out.flushTemp(r) + defer out.src.Seek(0, 0) + if err != nil { + return + } + + out.reader = csv.NewReader(r) + out.reader.ReuseRecord = true + + // Header + aux, err := out.reader.Read() + out.header = append(out.header, aux...) + if err != nil { + return + } + + out.row = make(map[string]string, len(out.header)) + + for { + _, err = out.reader.Read() + if err == io.EOF { + return out, nil + } else if err != nil { + return + } + + // Entry count + out.count++ + } +} + +// Cleanup should be called before we stop using the decoder +func (d *decoder) Cleanup() error { + return os.Remove(d.src.Name()) +} + +// Ident returns the assigned identifier +func (d *decoder) Ident() string { + return d.ident +} + +// Fields returns every available field in this dataset +func (d *decoder) Fields() []string { + return d.header +} + +// Reset resets the decoder to the start +func (d *decoder) Reset() error { + _, err := d.src.Seek(0, 0) + d.skipHead = true + return err +} + +// Next returns the field: value mapping for the next row +func (d *decoder) Next(out map[string]string) (more bool, err error) { + if !d.skipHead { + _, err = d.reader.Read() + if err != nil { + return + } + d.skipHead = true + } + + aux, err := d.reader.Read() + if err == io.EOF { + return false, nil + } else if err != nil { + return false, err + } + + for i, h := range d.header { + out[h] = aux[i] + } + + return true, nil +} + +// Count returns the total number of rows in the dataset +func (d *decoder) Count() uint64 { + return d.count +} + +func (d *decoder) flushTemp(r io.Reader) (_ io.Reader, err error) { + _, err = io.Copy(d.src, r) + if err != nil { + return + } + + d.src.Seek(0, 0) + return d.src, nil +} diff --git a/server/pkg/envoyx/csv/decoder_test.go b/server/pkg/envoyx/csv/decoder_test.go new file mode 100644 index 000000000..9f5a53a4f --- /dev/null +++ b/server/pkg/envoyx/csv/decoder_test.go @@ -0,0 +1,72 @@ +package csv + +import ( + "io" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDecoder(t *testing.T) { + req := require.New(t) + + t.Run("init & meta", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + req.Equal(uint64(3), dc.Count()) + }) + + t.Run("fields", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + hh := dc.Fields() + req.Contains(hh, "f1") + req.Contains(hh, "f2") + req.Contains(hh, "f3") + }) + + t.Run("iterate", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + aux := make(map[string]string) + var more bool + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r1f1", aux["f1"]) + req.Equal("r1f2", aux["f2"]) + req.Equal("r1f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r2f1", aux["f1"]) + req.Equal("r2f2", aux["f2"]) + req.Equal("r2f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r3f1", aux["f1"]) + req.Equal("r3f2", aux["f2"]) + req.Equal("r3f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.False(more) + }) +} + +func testReader() io.Reader { + src := `f1,f2,f3 +r1f1,r1f2,r1f3 +r2f1,r2f2,r2f3 +r3f1,r3f2,r3f3` + + return strings.NewReader(src) +} diff --git a/server/pkg/envoyx/json/decoder.go b/server/pkg/envoyx/json/decoder.go new file mode 100644 index 000000000..f4441f001 --- /dev/null +++ b/server/pkg/envoyx/json/decoder.go @@ -0,0 +1,145 @@ +package csv + +import ( + "encoding/json" + "io" + "io/ioutil" + "os" + "strings" + + "github.com/gabriel-vasile/mimetype" +) + +type ( + decoder struct { + ident string + + src *os.File + reader *json.Decoder + + header []string + count uint64 + } +) + +func CanDecodeFile(f io.Reader) bool { + m, err := mimetype.DetectReader(f) + if err != nil { + return false + } + + return CanDecodeExt(m.Extension()) +} + +func CanDecodeMime(m string) bool { + return m == "application/json" || m == "application/jsonlines" +} + +func CanDecodeExt(ext string) bool { + pt := strings.Split(ext, ".") + ext = strings.TrimSpace(pt[len(pt)-1]) + return ext == "jsonl" || ext == "json" || ext == "ndjson" +} + +// Decoder inits a new csv decoder from the given reader +// +// @todo hold small files in mem to avoid needles disc access +func Decoder(r io.Reader, ident string) (out *decoder, err error) { + out = &decoder{ + ident: ident, + } + + out.src, err = ioutil.TempFile(os.TempDir(), "*.ndjson") + if err != nil { + return + } + + r, err = out.flushTemp(r) + defer out.src.Seek(0, 0) + if err != nil { + return + } + + out.reader = json.NewDecoder(r) + + seenHeader := make(map[string]bool) + + var aux map[string]string + for out.reader.More() { + err = out.reader.Decode(&aux) + if err == io.EOF { + return out, nil + } else if err != nil { + return + } + + for f := range aux { + if seenHeader[f] { + continue + } + seenHeader[f] = true + + out.header = append(out.header, f) + } + + // Entry count + out.count++ + } + + return +} + +// Cleanup should be called before we stop using the decoder +func (d *decoder) Cleanup() error { + return os.Remove(d.src.Name()) +} + +// Ident returns the assigned identifier +func (d *decoder) Ident() string { + return d.ident +} + +// Fields returns every available field in this dataset +func (d *decoder) Fields() []string { + return d.header +} + +// Reset resets the decoder to the start +func (d *decoder) Reset() error { + _, err := d.src.Seek(0, 0) + return err +} + +// Next returns the field: value mapping for the next row +func (d *decoder) Next(out map[string]string) (more bool, err error) { + err = d.reader.Decode(&out) + if err == io.EOF { + return false, nil + } else if err != nil { + return false, err + } + + // Empty out missing fields to keep consistent with CSV + for _, h := range d.header { + if _, ok := out[h]; !ok { + out[h] = "" + } + } + + return true, nil +} + +// Count returns the total number of rows in the dataset +func (d *decoder) Count() uint64 { + return d.count +} + +func (d *decoder) flushTemp(r io.Reader) (_ io.Reader, err error) { + _, err = io.Copy(d.src, r) + if err != nil { + return + } + + d.src.Seek(0, 0) + return d.src, nil +} diff --git a/server/pkg/envoyx/json/decoder_test.go b/server/pkg/envoyx/json/decoder_test.go new file mode 100644 index 000000000..374032036 --- /dev/null +++ b/server/pkg/envoyx/json/decoder_test.go @@ -0,0 +1,71 @@ +package csv + +import ( + "io" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDecoder(t *testing.T) { + req := require.New(t) + + t.Run("init & meta", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + req.Equal(uint64(3), dc.Count()) + }) + + t.Run("fields", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + hh := dc.Fields() + req.Contains(hh, "f1") + req.Contains(hh, "f2") + req.Contains(hh, "f3") + }) + + t.Run("iterate", func(t *testing.T) { + dc, err := Decoder(testReader(), "test.csv") + req.NoError(err) + + aux := make(map[string]string) + var more bool + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r1f1", aux["f1"]) + req.Equal("r1f2", aux["f2"]) + req.Equal("r1f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r2f1", aux["f1"]) + req.Equal("r2f2", aux["f2"]) + req.Equal("r2f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.True(more) + req.Equal("r3f1", aux["f1"]) + req.Equal("r3f2", aux["f2"]) + req.Equal("r3f3", aux["f3"]) + + more, err = dc.Next(aux) + req.NoError(err) + req.False(more) + }) +} + +func testReader() io.Reader { + src := `{"f1": "r1f1", "f2": "r1f2", "f3": "r1f3"} +{"f1": "r2f1", "f2": "r2f2", "f3": "r2f3"} +{"f1": "r3f1", "f2": "r3f2", "f3": "r3f3"}` + + return strings.NewReader(src) +}