189 lines
3.1 KiB
Go
189 lines
3.1 KiB
Go
package cmd
|
||
|
||
import (
|
||
"encoding/csv"
|
||
"fmt"
|
||
"github.com/xuri/excelize/v2"
|
||
"io"
|
||
"log"
|
||
"math"
|
||
"os"
|
||
"regexp"
|
||
"strconv"
|
||
"time"
|
||
)
|
||
|
||
type (
|
||
Reader struct {
|
||
Path string
|
||
Index int
|
||
}
|
||
Writer struct {
|
||
File string
|
||
Limit int
|
||
}
|
||
Merge struct {
|
||
reader Reader
|
||
writer Writer
|
||
|
||
file *excelize.File
|
||
sw *excelize.StreamWriter
|
||
|
||
titles []interface{}
|
||
fileIndex int
|
||
total int
|
||
rowIndex int
|
||
}
|
||
)
|
||
|
||
func NewMerge(r Reader, w Writer) *Merge {
|
||
m := &Merge{
|
||
reader: r,
|
||
writer: w,
|
||
}
|
||
m.open()
|
||
return m
|
||
}
|
||
|
||
func (m *Merge) Merge() error {
|
||
begin := time.Now()
|
||
defer func() {
|
||
log.Printf("mergeCsvToExcel:耗时 %s\n", time.Now().Sub(begin).String())
|
||
if err := m.Save(); err != nil {
|
||
log.Println(err)
|
||
}
|
||
}()
|
||
|
||
for i := 0; i <= m.reader.Index; i++ {
|
||
filename := fmt.Sprintf("%s/data_%d_0.csv", m.reader.Path, i)
|
||
csvOpen, err := os.Open(filename)
|
||
|
||
if err != nil {
|
||
return fmt.Errorf("打开读取文件%s失败：%w", filename, err)
|
||
}
|
||
csvReader := csv.NewReader(csvOpen)
|
||
|
||
frist := true
|
||
for {
|
||
record, err := csvReader.Read()
|
||
if err == io.EOF {
|
||
break
|
||
} else if err != nil {
|
||
return fmt.Errorf("读取文件%s错误：%w", filename, err)
|
||
}
|
||
|
||
row := transform(record)
|
||
|
||
//不是第一个文件时，跳过第一条数据
|
||
if frist {
|
||
frist = false
|
||
|
||
if i == 0 {
|
||
m.WriteTitle(row)
|
||
}
|
||
continue
|
||
}
|
||
|
||
m.Write(row)
|
||
|
||
}
|
||
csvOpen.Close()
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func (m *Merge) WriteTitle(titles []interface{}) error {
|
||
if titles != nil {
|
||
m.titles = titles
|
||
}
|
||
if m.titles != nil {
|
||
return m.Write(m.titles)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (m *Merge) Write(values []interface{}) error {
|
||
|
||
cell, err := excelize.CoordinatesToCellName(1, m.rowIndex+1)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
err = m.sw.SetRow(cell, values)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
m.count()
|
||
return nil
|
||
}
|
||
|
||
func (m *Merge) reset() (err error) {
|
||
if m.file != nil {
|
||
if err := m.Save(); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
m.fileIndex++
|
||
m.rowIndex = 0
|
||
return m.open()
|
||
}
|
||
|
||
func (m *Merge) count() {
|
||
m.total++
|
||
m.rowIndex++
|
||
if m.rowIndex > m.writer.GetLimit() {
|
||
m.reset()
|
||
}
|
||
}
|
||
|
||
func (m *Merge) open() (err error) {
|
||
m.file = excelize.NewFile()
|
||
m.sw, err = m.file.NewStreamWriter("Sheet1")
|
||
m.WriteTitle(nil)
|
||
return err
|
||
}
|
||
|
||
func (m *Merge) Save() error {
|
||
//忽略只有标题的文件
|
||
if m.titles != nil && m.rowIndex == 1 {
|
||
return nil
|
||
}
|
||
|
||
if err := m.sw.Flush(); err != nil {
|
||
return err
|
||
}
|
||
|
||
return m.file.SaveAs(m.writer.GetFileName(m.fileIndex))
|
||
}
|
||
|
||
func (m *Merge) Clear() error {
|
||
return os.RemoveAll(m.reader.Path)
|
||
}
|
||
|
||
// GetFileName 获取文件名
|
||
func (w *Writer) GetFileName(fileIndex int) string {
|
||
ex := regexp.MustCompile("(\\..*)")
|
||
name := ex.ReplaceAllFunc([]byte(w.File), func(b []byte) []byte {
|
||
i := []byte("_" + strconv.Itoa(fileIndex))
|
||
ret := make([]byte, len(b)+len(i))
|
||
copy(ret, i)
|
||
copy(ret[len(i):], b)
|
||
return ret
|
||
})
|
||
return string(name)
|
||
}
|
||
|
||
func (w *Writer) GetLimit() int {
|
||
//excel 单表最大100w行数据
|
||
return int(math.Min(float64(w.Limit), 1000000))
|
||
}
|
||
|
||
func transform(record []string) []interface{} {
|
||
result := make([]interface{}, len(record))
|
||
for i2, s := range record {
|
||
result[i2] = s
|
||
}
|
||
return result
|
||
}
|