package service import ( "encoding/csv" "fmt" "github.com/360EntSecGroup-Skylar/excelize/v2" consts "github.com/easysoft/zendata/internal/pkg/const" logUtils "github.com/easysoft/zendata/pkg/utils/log" "github.com/easysoft/zendata/pkg/utils/vari" ) const ( sheetName = "Sheet1" ) var ( csvWriter *csv.Writer ) func (s *OutputService) GenExcel() { records := s.GenRecords() var f *excelize.File csvData := make([][]string, 0) if vari.GlobalVars.OutputFormat == consts.FormatExcel { f = excelize.NewFile() index := f.NewSheet(sheetName) f.SetActiveSheet(index) s.printExcelHeader(f) } else if vari.GlobalVars.OutputFormat == consts.FormatCsv { s.printCsvHeader(&csvData) csvWriter = csv.NewWriter(logUtils.OutputFileWriter) } for i, record := range records { csvRow := make([]string, 0) j := 0 for _, field := range vari.GlobalVars.ExportFields { val := record[field] if vari.GlobalVars.OutputFormat == consts.FormatExcel { colName, _ := excelize.CoordinatesToCellName(j+1, i+2) f.SetCellValue(sheetName, colName, val) } else if vari.GlobalVars.OutputFormat == consts.FormatCsv { csvRow = append(csvRow, fmt.Sprintf("%v", val)) } j++ } csvData = append(csvData, csvRow) } var err error if vari.GlobalVars.OutputFormat == consts.FormatExcel { err = f.SaveAs(logUtils.OutputFilePath) } else if vari.GlobalVars.OutputFormat == consts.FormatCsv { err = csvWriter.WriteAll(csvData) csvWriter.Flush() } if err != nil { logUtils.PrintErrMsg(err.Error()) } return } func (s *OutputService) printExcelHeader(f *excelize.File) { headerLine := "" for idx, field := range vari.GlobalVars.ExportFields { colName, _ := excelize.CoordinatesToCellName(idx+1, 1) f.SetCellValue(sheetName, colName, field) } logUtils.PrintLine(headerLine) } func (s *OutputService) printCsvHeader(csvData *[][]string) { csvRow := make([]string, 0) for _, field := range vari.GlobalVars.ExportFields { csvRow = append(csvRow, field) } *csvData = append(*csvData, csvRow) }