github.com/easysoft/zendata@v0.0.0-20240513203326-705bd5a7fd67/internal/pkg/service/output-excel.go (about)

     1  package service
     2  
     3  import (
     4  	"encoding/csv"
     5  	"fmt"
     6  
     7  	"github.com/360EntSecGroup-Skylar/excelize/v2"
     8  	consts "github.com/easysoft/zendata/internal/pkg/const"
     9  	logUtils "github.com/easysoft/zendata/pkg/utils/log"
    10  	"github.com/easysoft/zendata/pkg/utils/vari"
    11  )
    12  
    13  const (
    14  	sheetName = "Sheet1"
    15  )
    16  
    17  var (
    18  	csvWriter *csv.Writer
    19  )
    20  
    21  func (s *OutputService) GenExcel() {
    22  	records := s.GenRecords()
    23  
    24  	var f *excelize.File
    25  	csvData := make([][]string, 0)
    26  
    27  	if vari.GlobalVars.OutputFormat == consts.FormatExcel {
    28  		f = excelize.NewFile()
    29  		index := f.NewSheet(sheetName)
    30  		f.SetActiveSheet(index)
    31  
    32  		s.printExcelHeader(f)
    33  	} else if vari.GlobalVars.OutputFormat == consts.FormatCsv {
    34  		s.printCsvHeader(&csvData)
    35  		csvWriter = csv.NewWriter(logUtils.OutputFileWriter)
    36  	}
    37  
    38  	for i, record := range records {
    39  		csvRow := make([]string, 0)
    40  
    41  		j := 0
    42  		for _, field := range vari.GlobalVars.ExportFields {
    43  			val := record[field]
    44  
    45  			if vari.GlobalVars.OutputFormat == consts.FormatExcel {
    46  				colName, _ := excelize.CoordinatesToCellName(j+1, i+2)
    47  				f.SetCellValue(sheetName, colName, val)
    48  
    49  			} else if vari.GlobalVars.OutputFormat == consts.FormatCsv {
    50  				csvRow = append(csvRow, fmt.Sprintf("%v", val))
    51  			}
    52  
    53  			j++
    54  		}
    55  
    56  		csvData = append(csvData, csvRow)
    57  	}
    58  
    59  	var err error
    60  	if vari.GlobalVars.OutputFormat == consts.FormatExcel {
    61  		err = f.SaveAs(logUtils.OutputFilePath)
    62  	} else if vari.GlobalVars.OutputFormat == consts.FormatCsv {
    63  		err = csvWriter.WriteAll(csvData)
    64  		csvWriter.Flush()
    65  	}
    66  	if err != nil {
    67  		logUtils.PrintErrMsg(err.Error())
    68  	}
    69  
    70  	return
    71  }
    72  
    73  func (s *OutputService) printExcelHeader(f *excelize.File) {
    74  	headerLine := ""
    75  	for idx, field := range vari.GlobalVars.ExportFields {
    76  		colName, _ := excelize.CoordinatesToCellName(idx+1, 1)
    77  		f.SetCellValue(sheetName, colName, field)
    78  	}
    79  
    80  	logUtils.PrintLine(headerLine)
    81  }
    82  
    83  func (s *OutputService) printCsvHeader(csvData *[][]string) {
    84  	csvRow := make([]string, 0)
    85  
    86  	for _, field := range vari.GlobalVars.ExportFields {
    87  		csvRow = append(csvRow, field)
    88  	}
    89  
    90  	*csvData = append(*csvData, csvRow)
    91  }