Prepare JSON & CSV decoders for large imports
This commit is contained in:
Tomaž Jerman
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user