导出系统

internal/task/export_finalize.go

@@ -2,7 +2,10 @@ package task
 
 import (
 	"context"
+	"encoding/csv"
+	"errors"
 	"fmt"
+	"io"
 	"os"
 
 	"github.com/hibiken/asynq"
@@ -10,6 +13,7 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/break/junhong_cmp_fiber/internal/exporter"
+	"github.com/break/junhong_cmp_fiber/internal/model"
 	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
 	"github.com/break/junhong_cmp_fiber/pkg/constants"
 	"github.com/break/junhong_cmp_fiber/pkg/storage"
@@ -102,8 +106,7 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task *
 		return nil
 	}
 
-	strategy, ok := h.sceneRegistry.Get(exportTask.Scene)
-	if !ok {
+	if !h.sceneRegistry.IsSupported(exportTask.Scene) {
 		if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "导出场景未注册"); err != nil {
 			return err
 		}
@@ -115,23 +118,22 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task *
 		return err
 	}
 
-	rows := make([][]string, 0)
-	for _, shard := range shards {
-		if shard.Status != constants.ExportShardStatusSuccess {
-			continue
-		}
-		currentRows, err := strategy.QueryRows(ctx, exportTask, shard.CursorStart, shard.CursorEnd)
-		if err != nil {
-			return h.handleFinalizeError(ctx, exportTask.ID, updater, "读取分片数据失败", err)
-		}
-		rows = append(rows, currentRows...)
+	headers := exporter.ResolvedHeadersFromTask(exportTask)
+	if len(headers) == 0 {
+		return h.handleFinalizeError(ctx, exportTask.ID, updater, "读取导出表头失败", errors.New("resolved_headers 为空"))
 	}
 
-	localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows)
+	mergedCSVPath, cleanupCSV, err := h.mergeShardCSVFiles(ctx, headers, shards)
 	if err != nil {
 		return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
 	}
-	defer func() { _ = os.Remove(localPath) }()
+	defer cleanupCSV()
+
+	localPath, fileSize, cleanupFinal, err := prepareFinalExportFile(exportTask.Format, mergedCSVPath)
+	if err != nil {
+		return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
+	}
+	defer cleanupFinal()
 
 	fileName := fmt.Sprintf("%s.%s", exportTask.TaskNo, exportTask.Format)
 	fileKey, err := uploadExportFile(ctx, h.storageSvc, localPath, fileName)
@@ -146,6 +148,104 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task *
 	return nil
 }
 
+func (h *ExportFinalizeHandler) mergeShardCSVFiles(ctx context.Context, headers []string, shards []*model.ExportShardTask) (string, func(), error) {
+	if h.storageSvc == nil || h.storageSvc.Provider() == nil {
+		return "", nil, errors.New("对象存储服务未配置")
+	}
+
+	file, err := os.CreateTemp("", "export-merged-*.csv")
+	if err != nil {
+		return "", nil, err
+	}
+
+	cleanup := func() {
+		_ = os.Remove(file.Name())
+	}
+
+	writer := csv.NewWriter(file)
+	if err := writer.Write(headers); err != nil {
+		file.Close()
+		cleanup()
+		return "", nil, err
+	}
+
+	for _, shard := range shards {
+		if shard.Status != constants.ExportShardStatusSuccess {
+			continue
+		}
+		if shard.FileKey == "" {
+			file.Close()
+			cleanup()
+			return "", nil, fmt.Errorf("分片 %d 缺少文件Key", shard.ShardNo)
+		}
+		if err := h.appendShardFile(ctx, writer, shard.FileKey); err != nil {
+			file.Close()
+			cleanup()
+			return "", nil, err
+		}
+	}
+
+	writer.Flush()
+	if err := writer.Error(); err != nil {
+		file.Close()
+		cleanup()
+		return "", nil, err
+	}
+	if err := file.Sync(); err != nil {
+		file.Close()
+		cleanup()
+		return "", nil, err
+	}
+	if err := file.Close(); err != nil {
+		cleanup()
+		return "", nil, err
+	}
+
+	return file.Name(), cleanup, nil
+}
+
+func (h *ExportFinalizeHandler) appendShardFile(ctx context.Context, writer *csv.Writer, key string) error {
+	reader, err := h.storageSvc.Provider().Download(ctx, key)
+	if err != nil {
+		return err
+	}
+	defer reader.Close()
+
+	csvReader := csv.NewReader(reader)
+	for {
+		row, err := csvReader.Read()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return err
+		}
+		if err := writer.Write(row); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func prepareFinalExportFile(format, mergedCSVPath string) (string, int64, func(), error) {
+	switch format {
+	case constants.ExportTaskFormatCSV:
+		stat, err := os.Stat(mergedCSVPath)
+		if err != nil {
+			return "", 0, nil, err
+		}
+		return mergedCSVPath, stat.Size(), func() {}, nil
+	case constants.ExportTaskFormatXLSX:
+		xlsxPath, fileSize, err := writeXLSXFromCSV(mergedCSVPath)
+		if err != nil {
+			return "", 0, nil, err
+		}
+		return xlsxPath, fileSize, func() { _ = os.Remove(xlsxPath) }, nil
+	default:
+		return "", 0, nil, fmt.Errorf("不支持的导出格式: %s", format)
+	}
+}
+
 func (h *ExportFinalizeHandler) handleFinalizeError(ctx context.Context, taskID uint, updater uint, prefix string, err error) error {
 	if isFinalRetry(ctx) {
 		markErr := h.taskStore.MarkFailed(ctx, taskID, updater, prefix+": "+err.Error())