diff --git a/cmd/cmd/csv.go b/cmd/cmd/csv.go new file mode 100644 index 0000000..5e4ebfb --- /dev/null +++ b/cmd/cmd/csv.go @@ -0,0 +1,220 @@ +package cmd + +import ( + "encoding/csv" + "excel_export/biz/config" + "excel_export/biz/db" + "excel_export/biz/export" + "fmt" + "github.com/xuri/excelize/v2" + "io" + "log" + "os" + "strconv" + "sync" + "time" +) + +type Csv struct { + conf *config.Config + dirTemp string +} + +func NewCsv(conf *config.Config) *Csv { + return &Csv{ + conf: conf, + } +} + +func (e *Csv) Export(sysName, jobName string, begin, end time.Time, batch int) error { + job, dbStr, err := config.GetJob(e.conf, sysName, jobName) + if err != nil { + return err + } + + d, err := db.NewDb(dbStr) + if err != nil { + return err + } + + return e.JobHandler(job, d, map[string]interface{}{ + "begin": begin, + "end": end, + "last": 0, + }, batch) +} + +func (e *Csv) JobHandler(job config.Job, d export.DataFetcher, params map[string]interface{}, batch int) error { + + for i, task := range job.Tasks { + fmt.Printf("执行导出任务:%d\n", i+1) + if err := e.TaskExport(d, task, params, batch); err != nil { + return err + } + } + return nil +} + +func (e *Csv) TaskExport(d export.DataFetcher, t config.Task, params map[string]interface{}, batch int) error { + var i int + var wg sync.WaitGroup + for i = 0; i < 1000; i++ { + + f, err := e.getCsvFile(i) + if err != nil { + return err + } + + sql := t.GetSql(params) + + e := export.NewCsvExporter(d, f) + e.(*export.CsvExporter).WaitGroup(&wg) + wg.Add(1) + e.Export(sql+" limit "+strconv.Itoa(batch), t.PK) + + count, last := e.Last() + fmt.Printf("已导出 %d 条数据\n", batch*i+count) + + if count == 0 { + return nil + } + + if count < batch { + break + } + + params["last"] = last + time.Sleep(time.Microsecond * 30) + } + wg.Wait() + fmt.Println("tempDir", e.dirTemp) + //todo 合并csv文件,并删除 临时目录 + err := e.mergeCsvToExcel(e.dirTemp, i) + fmt.Println(err) + + //重置临时路径 + e.dirTemp = "" + + return nil +} + +func (e *Csv) getCsvFile(index int) (*export.Csv, error) { + + if e.dirTemp == "" { + path, err := os.MkdirTemp(os.TempDir(), "") + if err != nil { + return nil, err + } + e.dirTemp = path + } + + filename := e.dirTemp + "/data_{index}.csv" + + f := export.NewCsv(filename, map[string]string{ + "index": strconv.Itoa(index), + }) + + return f, nil +} + +func (e *Csv) mergeCsv(path string, max int) error { + csvFile, err := os.Create("aa.csv") + if err != nil { + return fmt.Errorf("打开写入文件失败:%w", err) + } + defer csvFile.Close() + csvWriter := csv.NewWriter(csvFile) + + for i := 0; i <= max; i++ { + filename := fmt.Sprintf("%s/data_%d_0.csv", path, i) + csvOpen, err := os.Open(filename) + + if err != nil { + return fmt.Errorf("打开读取文件%s失败:%w", filename, err) + } + csvReader := csv.NewReader(csvOpen) + var index int + for { + record, err := csvReader.Read() + if err == io.EOF { + break + } else if err != nil { + return fmt.Errorf("读取文件%s错误:%w", filename, err) + } + + index++ + //不是第一个文件时,跳过第一条数据 + if index == 1 && i != 0 { + //continue + } + fmt.Println(record) + csvWriter.Write(record) + + } + + csvOpen.Close() + } + csvWriter.Flush() + return nil +} + +func (e *Csv) mergeCsvToExcel(path string, max int) error { + f := excelize.NewFile() + defer func() { + if err := f.Close(); err != nil { + log.Println(err) + } + }() + + sheet, err := f.NewStreamWriter("Sheet1") + if err != nil { + log.Println(err) + return err + } + + var index int + for i := 0; i <= max; i++ { + filename := fmt.Sprintf("%s/data_%d_0.csv", 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) + } + + //不是第一个文件时,跳过第一条数据 + if frist && i != 0 { + frist = false + continue + } + + index++ + + cell, _ := excelize.CoordinatesToCellName(1, index) + + rec := make([]interface{}, len(record)) + for i2, s := range record { + rec[i2] = s + } + + sheet.SetRow(cell, rec) + + } + csvOpen.Close() + } + + if err := sheet.Flush(); err != nil { + return err + } + + return f.SaveAs("ss.xlsx") +} diff --git a/cmd/cmd/csv_test.go b/cmd/cmd/csv_test.go new file mode 100644 index 0000000..755a063 --- /dev/null +++ b/cmd/cmd/csv_test.go @@ -0,0 +1,24 @@ +package cmd + +import ( + "fmt" + "github.com/stretchr/testify/assert" + "os" + "testing" +) + +func TestCsv_getCsvFile(t *testing.T) { + + c := &Csv{} + f, err := c.getCsvFile(1) + assert.Nil(t, err) + assert.NotNil(t, f) + f.Open() + f.Close() + fmt.Println(c.dirTemp) + + assert.DirExists(t, c.dirTemp) + os.RemoveAll(c.dirTemp) + assert.NoDirExists(t, c.dirTemp) + +}