l_excel_export/export.go

package excel_export

import (
	"encoding/json"
	"fmt"
	"gitea.cdlsxd.cn/self-tools/l_excel_export/pkg"
	"gitea.cdlsxd.cn/self-tools/l_excel_export/types"
	"time"

	"gitea.cdlsxd.cn/self-tools/l_excel_export/export_err"
	"github.com/xuri/excelize/v2"
	"os"
)

type ExportExcel struct {
	config    *Config
	password  string
	saveFunc  func(file *excelize.File) (url string, err error)
	sheetName string
	logPath   string
	jobName   string
	task      *types.Task

	exportData [][]interface{}
	header     []interface{}
}

func NewExport(config *Config, opts ...Option) (*ExportExcel, error) {
	export := &ExportExcel{
		config: config,
	}
	for _, opt := range opts {
		opt(export) // 应用选项
	}
	err := export.check()
	return export, err
}

func (e *ExportExcel) Run() (taskId string, err error) {
	err = e.init()
	if err != nil {
		return
	}
	f, err := e.getFile()
	if err != nil {
		return
	}
	go func(f *excelize.File) {
		defer func() {
			if r := recover(); r != nil {
				f.Close()
				fmt.Println(r)
			}
		}()
		e.task.Status = Running
		e.updateTask()
		err = e.run(f)
		if err != nil {
			e.task.Status = Err
			e.updateTask()
			return
		}
		e.task.Ftime = time.Now().Format(time.DateTime)
		e.finish(f)
	}(f)
	return e.task.TaskId, nil
}

func (e *ExportExcel) run(f *excelize.File) (err error) {

	index, err := f.NewStreamWriter(e.getSheetName())
	if err != nil {
		return err
	}
	index.SetRow("A1", e.header)
	if err != nil {
		return err
	}
	count := len(e.exportData)
	for i, v := range e.exportData {
		cell, _ := excelize.CoordinatesToCellName(1, i+2)
		index.SetRow(cell, v)
		e.task.Process = i * 100 / count
		e.updateTask()
	}

	if err = e.save(f); err != nil {
		return err
	}
	return err
}

func (e *ExportExcel) finish(f *excelize.File) {
	f.Close()
	e.task.Status = Finish
	e.task.Process = 100
	e.updateTask()
}

func (e *ExportExcel) save(f *excelize.File) error {
	e.task.Url = e.getUrl()
	return f.SaveAs(e.task.Url)
}

func (e *ExportExcel) getUrl() string {
	return fmt.Sprintf("%s/%s_%s%s", e.config.SavePath, e.config.FileName, e.task.TaskId, e.config.Ext)
}

func (e *ExportExcel) init() error {
	e.task = &types.Task{
		TaskId:   pkg.CreateTaskId(),
		Process:  0,
		Url:      "",
		Ctime:    time.Now().Format(time.DateTime),
		Ftime:    "",
		Status:   Init,
		FileAddr: "",
	}
	err := e.updateTask()
	if err != nil {
		return err
	}
	dataMap := pkg.GetData(e.config.Data)
	if len(dataMap) == 0 {
		return export_err.ErrDataError
	}

	if len(e.config.Head) == 0 {
		return export_err.ErrHeadNotSet
	}

	for _, v := range e.config.Head {

		e.header = append(e.header, v.ColName)
	}

	//todo:这一步目的是为了保证数据与header未知一致，但是在大数据量的时候存在性能问题
	for _, v := range dataMap {
		var (
			slice []interface{}
		)
		for _, vv := range e.config.Head {
			slice = append(slice, v[vv.FieldName])
		}
		e.exportData = append(e.exportData, slice)
	}
	//检测文件是否存在
	if _, err = os.Stat(e.config.SavePath); os.IsNotExist(err) {
		// 文件夹不存在，尝试创建
		err = os.MkdirAll(e.config.SavePath, os.ModePerm)
		if err != nil {
			return export_err.ErrCreateExcelPathFail
		}
	} else if err != nil {
		return export_err.ErrCheckExcelPathFail
	}

	return nil
}

func (e *ExportExcel) updateTask() error {

	file, err := os.OpenFile(e.logFile(), os.O_RDWR|os.O_TRUNC|os.O_CREATE, 0766)
	if err != nil {
		return err
	}
	defer file.Close()

	taskInfo, err := json.Marshal(e.task)
	if err != nil {
		return err
	}
	jsonInfo := string(taskInfo)
	_, err = file.WriteString(jsonInfo)
	if err != nil {
		return err
	}
	return nil
}

func (e *ExportExcel) logFile() string {
	return fmt.Sprintf("%s/%s", e.logPath, e.task.TaskId)
}

func (e *ExportExcel) getFile() (*excelize.File, error) {
	f, err := excelize.OpenFile(e.path())
	if err != nil {
		f = excelize.NewFile()
	}
	return f, nil
}

func (e *ExportExcel) path() string {
	return fmt.Sprintf("%s/%s.%s", e.config.SavePath, e.config.FileName, e.config.Ext)
}

func (e *ExportExcel) check() (err error) {
	if len(e.jobName) == 0 {
		e.jobName = time.Now().Format("default")
	}
	if e.config.SavePath == "" && e.saveFunc == nil {
		return export_err.ErrNotSetSaveWay
	}

	if len(e.logPath) == 0 {
		e.logPath, err = pkg.DefaultLogPath(e.jobName)
		return
	}
	return
}

func (e *ExportExcel) getSheetName() string {
	if len(e.sheetName) == 0 {
		e.sheetName = "Sheet1"
	}
	return e.sheetName
}