导出系统

internal/task/export_common.go

@@ -5,6 +5,7 @@ import (
 	"encoding/csv"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"time"
@@ -82,18 +83,18 @@ func isFinalRetry(ctx context.Context) bool {
 	return retryCount >= maxRetry
 }
 
-func writeExportFile(format string, headers []string, rows [][]string) (string, int64, error) {
+func writeExportFile(format string, headers []string, rows [][]string, withoutHeader bool) (string, int64, error) {
 	switch format {
 	case constants.ExportTaskFormatCSV:
-		return writeCSVFile(headers, rows)
+		return writeCSVFile(headers, rows, withoutHeader)
 	case constants.ExportTaskFormatXLSX:
-		return writeXLSXFile(headers, rows)
+		return writeXLSXFile(headers, rows, withoutHeader)
 	default:
 		return "", 0, fmt.Errorf("不支持的导出格式: %s", format)
 	}
 }
 
-func writeCSVFile(headers []string, rows [][]string) (string, int64, error) {
+func writeCSVFile(headers []string, rows [][]string, withoutHeader bool) (string, int64, error) {
 	file, err := os.CreateTemp("", "export-*.csv")
 	if err != nil {
 		return "", 0, err
@@ -101,8 +102,10 @@ func writeCSVFile(headers []string, rows [][]string) (string, int64, error) {
 	defer file.Close()
 
 	writer := csv.NewWriter(file)
-	if err := writer.Write(headers); err != nil {
-		return "", 0, err
+	if !withoutHeader {
+		if err := writer.Write(headers); err != nil {
+			return "", 0, err
+		}
 	}
 	for _, row := range rows {
 		if err := writer.Write(row); err != nil {
@@ -125,7 +128,7 @@ func writeCSVFile(headers []string, rows [][]string) (string, int64, error) {
 	return file.Name(), stat.Size(), nil
 }
 
-func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) {
+func writeXLSXFile(headers []string, rows [][]string, withoutHeader bool) (string, int64, error) {
 	file := excelize.NewFile()
 	sheet := file.GetSheetName(file.GetActiveSheetIndex())
 
@@ -134,16 +137,20 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) {
 		return "", 0, err
 	}
 
-	headerValues := make([]interface{}, 0, len(headers))
-	for _, item := range headers {
-		headerValues = append(headerValues, item)
-	}
-	if err := streamWriter.SetRow("A1", headerValues); err != nil {
-		return "", 0, err
+	rowIndex := 1
+	if !withoutHeader {
+		headerValues := make([]interface{}, 0, len(headers))
+		for _, item := range headers {
+			headerValues = append(headerValues, item)
+		}
+		if err := streamWriter.SetRow("A1", headerValues); err != nil {
+			return "", 0, err
+		}
+		rowIndex++
 	}
 
-	for i, row := range rows {
-		axis, err := excelize.CoordinatesToCellName(1, i+2)
+	for _, row := range rows {
+		axis, err := excelize.CoordinatesToCellName(1, rowIndex)
 		if err != nil {
 			return "", 0, err
 		}
@@ -154,6 +161,7 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) {
 		if err := streamWriter.SetRow(axis, values); err != nil {
 			return "", 0, err
 		}
+		rowIndex++
 	}
 
 	if err := streamWriter.Flush(); err != nil {
@@ -175,6 +183,90 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) {
 	return path, stat.Size(), nil
 }
 
+func writeXLSXFromCSV(csvPath string) (string, int64, error) {
+	source, err := os.Open(csvPath)
+	if err != nil {
+		return "", 0, err
+	}
+	defer source.Close()
+
+	file := excelize.NewFile()
+	sheet := file.GetSheetName(file.GetActiveSheetIndex())
+	streamWriter, err := file.NewStreamWriter(sheet)
+	if err != nil {
+		return "", 0, err
+	}
+
+	reader := csv.NewReader(source)
+	rowIndex := 1
+	for {
+		row, err := reader.Read()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return "", 0, err
+		}
+
+		values := make([]interface{}, 0, len(row))
+		for _, item := range row {
+			values = append(values, item)
+		}
+		axis, err := excelize.CoordinatesToCellName(1, rowIndex)
+		if err != nil {
+			return "", 0, err
+		}
+		if err := streamWriter.SetRow(axis, values); err != nil {
+			return "", 0, err
+		}
+		rowIndex++
+	}
+
+	if err := streamWriter.Flush(); err != nil {
+		return "", 0, err
+	}
+
+	path := filepath.Join(os.TempDir(), fmt.Sprintf("export-%d.xlsx", time.Now().UnixNano()))
+	if err := file.SaveAs(path); err != nil {
+		return "", 0, err
+	}
+	if err := file.Close(); err != nil {
+		return "", 0, err
+	}
+
+	stat, err := os.Stat(path)
+	if err != nil {
+		return "", 0, err
+	}
+	return path, stat.Size(), nil
+}
+
+func alignRowsToHeaders(rows [][]string, headers []string) [][]string {
+	if len(headers) == 0 {
+		return rows
+	}
+
+	aligned := make([][]string, 0, len(rows))
+	for _, row := range rows {
+		current := make([]string, len(headers))
+		copy(current, row)
+		aligned = append(aligned, current)
+	}
+	return aligned
+}
+
+func hasMisalignedRows(rows [][]string, headers []string) bool {
+	if len(headers) == 0 {
+		return false
+	}
+	for _, row := range rows {
+		if len(row) != len(headers) {
+			return true
+		}
+	}
+	return false
+}
+
 func uploadExportFile(ctx context.Context, storageSvc *storage.Service, localPath, fileName string) (string, error) {
 	if storageSvc == nil || storageSvc.Provider() == nil {
 		return "", errors.New("对象存储服务未配置")