featurebase/idk/csv/source.go

package csv

import (
	"context"
	"encoding/csv"
	"io"
	"net/http"
	"os"
	"reflect"
	"strings"
	"sync"

	"github.com/featurebasedb/featurebase/v3/idk"
	"github.com/featurebasedb/featurebase/v3/logger"
	"github.com/pkg/errors"
)

type Source struct {
	Files        chan string
	Header       []string
	IgnoreHeader bool
	JustDoIt     bool
	Log          logger.Logger

	schemaLock sync.Mutex
	schema     []idk.Field

	records      chan Record
	once         *sync.Once
	expectHeader bool
}

func (s *Source) Record() (idk.Record, error) {
	s.once.Do(func() { go s.run() })

	rec, ok := <-s.records
	if !ok {
		return nil, io.EOF
	}
	return rec, rec.err
}

type Record struct {
	data []interface{}
	err  error
}

func (r Record) Data() []interface{} {
	return r.data
}
func (r Record) Commit(ctx context.Context) error { return nil } // TODO do

func (r Record) Schema() interface{} { return nil }

func (s *Source) Schema() []idk.Field {
	s.schemaLock.Lock()
	defer s.schemaLock.Unlock()
	return s.schema
}

func (s *Source) run() {
	defer close(s.records)

	s.expectHeader = len(s.Header) == 0
	if !s.expectHeader {
		var err error
		s.schemaLock.Lock()
		s.schema, err = s.processHeader(s.Header)
		s.schemaLock.Unlock()
		if err != nil {
			s.records <- Record{err: errors.Wrapf(err, "processing given header: %+v", s.Header)}
			return
		}
	}

	for filename := range s.Files {
		s.processFile(filename)
	}
}

func (s *Source) processHeader(header []string) (schema []idk.Field, err error) {
	schema = make([]idk.Field, len(header))
	for i, val := range header {
		schema[i], err = idk.HeaderToField(val, s.Log)
		if err != nil && s.JustDoIt {
			schema[i] = idk.StringField{NameVal: strings.ToLower(val)}
		} else if err != nil {
			return nil, err
		}
	}
	return schema, nil
}

func (s *Source) processFile(name string) {
	s.Log.Printf("processFile: %s", name)

	f, err := openFileOrURL(name)
	if err != nil {
		s.records <- Record{err: errors.Wrapf(err, "opening %s", name)}
		return
	}
	defer f.Close()

	reader := csv.NewReader(f)
	reader.ReuseRecord = true
	reader.FieldsPerRecord = 0
	var nextErr error
	if s.expectHeader || s.IgnoreHeader {
		header, err := reader.Read()
		if err != nil {
			s.records <- Record{err: errors.Wrapf(err, "reading header from '%s'", name)}
			return
		}
		if s.expectHeader {
			newschema, err := s.processHeader(header)
			if err != nil {
				s.Log.Printf("processHeader error: %v\n", err)
				s.records <- Record{err: errors.Wrapf(err, "processing header from '%s': %+v", name, header)}
				return
			}
			if !reflect.DeepEqual(newschema, s.schema) {
				s.schema = newschema
				nextErr = idk.ErrSchemaChange
			}
		}
	}
	// weird hack to reuse memory rather than allocating eery
	// time. Would probably be better to get rid of the channel and
	// just the next CSV row on demand.
	recs := [2]Record{
		{data: make([]interface{}, len(s.schema)), err: nextErr},
		{data: make([]interface{}, len(s.schema))},
	}
	i := -1
	extraColumnsCount := 0
	row, err := reader.Read()
	for ; err == nil; row, err = reader.Read() {
		i += 1
		for j, val := range row {
			if len(s.schema) <= j {
				if extraColumnsCount == 0 {
					s.Log.Warnf("'%s': ignoring additional column(s) not included in the header specification", name)
				}
				extraColumnsCount++
				break
			}
			recs[i%2].data[j] = val
		}
		s.records <- recs[i%2]
		recs[i%2].err = nil
	}
	if extraColumnsCount > 0 {
		s.Log.Printf("Processing '%s': %d rows have more columns than header specification", name, extraColumnsCount)
	}
	if err != io.EOF {
		s.Log.Printf("ERROR Processing '%s': '%v'. Skipping rest of file.", name, err)
	}
}

func NewSource() *Source {
	return &Source{
		records: make(chan Record, 0), // nolint: gosimple // do not change buffer size!
	}
}

func openFileOrURL(name string) (io.ReadCloser, error) {
	var content io.ReadCloser
	if strings.HasPrefix(name, "http") {
		resp, err := http.Get(name)
		if err != nil {
			return nil, errors.Wrap(err, "getting via http")
		}
		if resp.StatusCode > 299 {
			return nil, errors.Errorf("got status %d via http.Get", resp.StatusCode)
		}
		content = resp.Body
	} else {
		f, err := os.Open(name)
		if err != nil {
			return nil, errors.Wrap(err, "opening file")
		}
		content = f
	}
	return content, nil
}

func (s *Source) Close() error {
	return nil
}