featurebase/idk/csv/cmd.go

package csv

import (
	"encoding/csv"
	"os"
	"path/filepath"
	"sync"

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

type Main struct {
	idk.Main     `flag:"!embed"`
	Files        []string `help:"List of files, URLs, or directories to ingest."`
	Header       []string `help:"Optional header. If not passed, first line of each file is used."`
	IgnoreHeader bool     `short:"g" help:"Ignore header in file and use configured header. You *must* configure a header."`
	JustDoIt     bool     `short:"j" help:"Any header field not in the appropriate format, just downcase, use it as the name and process the value as a String/set field"`

	files   chan string
	visited map[string]struct{}
}

func NewMain() *Main {
	m := &Main{
		Main:    *idk.NewMain(),
		files:   make(chan string),
		visited: make(map[string]struct{}),
	}
	m.Main.Namespace = "ingester_csv"

	once := &sync.Once{}
	m.NewSource = func() (idk.Source, error) {
		if len(m.Files) == 0 {
			return nil, errors.New("must provide at least one file or directory with --files")
		}
		once.Do(func() { go m.streamFileNames() })

		source := NewSource()
		source.Files = m.files
		source.Header = m.Header
		source.Log = m.Main.Log()
		source.once = &sync.Once{}
		source.JustDoIt = m.JustDoIt
		source.IgnoreHeader = m.IgnoreHeader
		return source, nil
	}
	return m
}

// ValidateHeaders is used during --dry-run to check that all headers provided at CLI or in
// CSV files can be parsed properly.
// Also returns the field structs as it would be parsed in a normal run. In the case of
// multiple CSV files, only returns the fields for the last file.
func (m *Main) ValidateHeaders() ([]idk.Field, error) {
	var err error
	fields := make([]idk.Field, 0)
	if len(m.Header) > 0 {
		fields, err = m.validateHeader(m.Header)
		if err != nil {
			return fields, errors.Wrap(err, "validating header")
		}
	} else {
		for _, f := range m.Files {
			f, err := openFileOrURL(f)
			if err != nil {
				return fields, errors.Wrapf(err, "opening file (%s)", f)
			}
			defer f.Close()
			reader := csv.NewReader(f)
			header, err := reader.Read()
			if err != nil {
				return fields, errors.Wrapf(err, "reading CSV file (%s)", f)
			}
			fields, err = m.validateHeader(header)
			if err != nil {
				return fields, errors.Wrap(err, "validating header")
			}
		}
	}
	return fields, nil
}

func (m *Main) validateHeader(h []string) ([]idk.Field, error) {
	fields := make([]idk.Field, 0)
	for _, val := range h {
		field, err := idk.HeaderToField(val, m.Log())
		if err != nil {
			return fields, errors.Wrapf(err, "invalid header (%v)", val)
		}
		fields = append(fields, field)
	}
	return fields, nil
}

func (m *Main) streamFileNames() {
	defer close(m.files)

	for _, filename := range m.Files {
		if err := filepath.Walk(filename, m.traverse); err != nil {
			m.Log().Printf("filepath.Walk(%s) %v", filename, err)
			continue
		}
	}
}

func (m *Main) traverse(path string, fi os.FileInfo, err error) error {
	if err != nil {
		m.Log().Printf("prevent panic by handling failure accessing a path %q: %v\n", path, err)
		return err
	}

	switch mode := fi.Mode(); {
	case mode.IsRegular():
		if _, ok := m.visited[path]; !ok {
			m.visited[path] = struct{}{}
			m.files <- path
		}

	case mode.IsDir():
		if _, ok := m.visited[path]; !ok {
			m.visited[path] = struct{}{}
		}

	case mode&os.ModeSymlink != 0:
		lnk, err := filepath.EvalSymlinks(path)
		if err != nil {
			return err
		}
		if _, ok := m.visited[lnk]; !ok {
			if err := filepath.Walk(lnk, m.traverse); err != nil {
				return err
			}
		}
	}

	return nil
}