cgimport/reader/csvreader.go

package reader

import (
	"fmt"
	"io"
	"strconv"
	"strings"

	"git.wecise.com/wecise/odb-go/schema"
	"github.com/spf13/cast"
	"github.com/wecisecode/util/merrs"
)

type CSVBlockReader struct {
	*LineReader
	schema    *schema.Schema
	classname string
	csvkeys   []string
}

func NewCSVBlockReader(filename string, classname string, reader io.Reader, schema *schema.Schema) *CSVBlockReader {
	return &CSVBlockReader{
		LineReader: NewLineReader(filename, reader),
		schema:     schema,
		classname:  classname,
	}
}

func (br *CSVBlockReader) ReadBlock(skiplines int) (block map[string]any, line string, linecount int, err error) {
	classname := br.classname
	ci := br.schema.GetClassInfo(classname)
	eof := false
	for {
		line, linecount, eof, err = br.ReadLine()
		if err != nil {
			return
		}
		if linecount == 1 {
			br.csvkeys = strings.Split(line, "^")
			line, linecount, eof, err = br.ReadLine()
			if err != nil {
				return
			}
		}
		if line == "" {
			if eof {
				return
			}
			continue
		}
		if linecount <= skiplines {
			continue
		}
		values := strings.Split(line, "^")
		if len(values) != len(br.csvkeys) {
			err = merrs.NewError(fmt.Sprint(br.filename, " format error, values count not match keys count, line ", br.linecount))
			return
		}
		block = map[string]any{}
		for i, k := range br.csvkeys {
			v := values[i]
			if v != "" {
				n := cast.ToInt64(v)
				if n != 0 || v == "0" {
					block[k] = n
					continue
				}
				f := cast.ToFloat64(v)
				if f != 0 || v == "0" || v == "0.0" || v == ".0" || v == "0." {
					block[k] = f
					continue
				}
				b := cast.ToBool(v)
				if v == "true" || v == "false" {
					block[k] = b
					continue
				}
				s, e := strconv.Unquote(v)
				if e == nil {
					v = s
				}
			}
			if ci != nil {
				fi := ci.DatakeyFieldinfos[k]
				if fi != nil {
					switch fi.Fieldtype {
					case "set<varchar>":
						s := v
						if strings.HasPrefix(s, "[") && strings.HasSuffix(s, "]") {
							s = s[1 : len(s)-1]
						}
						ss := cast.ToStringSlice(s)
						block[k] = ss
						continue
					case "varchar":
						block[k] = cast.ToString(v)
						continue
					case "bigint":
						block[k] = cast.ToInt64(v)
						continue
					}
				}
			}
			block[k] = v
		}
		return
	}
}