corteza/pkg/envoy/csv/decoder.go

package csv

import (
	"bytes"
	"context"
	"encoding/csv"
	"io"
	"strings"

	"github.com/cortezaproject/corteza-server/pkg/envoy"
	"github.com/cortezaproject/corteza-server/pkg/envoy/resource"
	"github.com/gabriel-vasile/mimetype"
)

type (
	// wrapper struct for csv related methods
	decoder struct{}

	// csv decoder wrapper for additional bits
	reader struct {
		c      *csv.Reader
		header []string
		count  uint64

		buff bytes.Buffer
	}
)

// Decoder initializes and returns a fresh CSV decoder
func Decoder() *decoder {
	return &decoder{}
}

// CanDecodeFile determines if the file can be determined by this decoder
func (y *decoder) CanDecodeFile(f io.Reader) bool {
	m, err := mimetype.DetectReader(f)
	if err != nil {
		return false
	}

	return y.CanDecodeExt(m.Extension())
}

func (y *decoder) CanDecodeMime(m string) bool {
	return m == "text/csv"
}

func (y *decoder) CanDecodeExt(ext string) bool {
	pt := strings.Split(ext, ".")
	return strings.TrimSpace(pt[len(pt)-1]) == "csv"
}

// Decode decodes the given io.Reader into a generic resource dataset
func (c *decoder) Decode(ctx context.Context, r io.Reader, do *envoy.DecoderOpts) ([]resource.Interface, error) {
	cr := &reader{}
	var buff bytes.Buffer

	// So we can reset to the start of the reader
	err := cr.prepare(io.TeeReader(r, &buff))
	if err != nil {
		return nil, err
	}

	cr.c = csv.NewReader(io.TeeReader(&buff, &cr.buff))

	// The first one is a header, so let's just get rid of it
	_, err = cr.c.Read()
	if err != nil {
		return nil, err
	}

	return []resource.Interface{resource.NewResourceDataset(do.Name, cr)}, nil
}

func (cr *reader) prepare(r io.Reader) (err error) {
	cReader := csv.NewReader(r)

	// Header
	cr.header, err = cReader.Read()
	if err != nil {
		return err
	}

	// Entry count
	for {
		_, err := cReader.Read()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}

		cr.count++
	}

	return nil
}

// Fields returns every available field in this dataset
func (cr *reader) Fields() []string {
	return cr.header
}

func (cr *reader) Reset() error {
	if len(cr.buff.Bytes()) == 0 {
		return nil
	}

	buff := cr.buff
	cr.buff = bytes.Buffer{}

	cr.c = csv.NewReader(io.TeeReader(&buff, &cr.buff))
	// The first one is a header, so let's just get rid of it
	_, err := cr.c.Read()
	if err != nil {
		return err
	}

	return nil
}

// Next returns the field: value mapping for the next row
func (cr *reader) Next() (map[string]string, error) {
	mr := make(map[string]string)
	rr, err := cr.c.Read()
	if err == io.EOF {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	for i, h := range cr.header {
		mr[h] = rr[i]
	}

	return mr, nil
}

func (cr *reader) Count() uint64 {
	return cr.count
}