导出系统

2026-06-15 15:28:29 +08:00
parent 2f0b24ce88
commit 7ec84fbc0f
17 changed files with 939 additions and 260 deletions
--- a/internal/task/export_common.go
+++ b/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("对象存储服务未配置")
--- a/internal/task/export_dispatch.go
+++ b/internal/task/export_dispatch.go
@@ -119,14 +119,22 @@ func (h *ExportDispatchHandler) HandleExportDispatch(ctx context.Context, task *
 		return asynq.SkipRetry
 	}

-	shards, totalRows, err := h.buildShards(ctx, exportTask, strategy, updater)
+	params := exporter.ParseExportParams(exportTask)
+	headers, err := strategy.Headers(ctx, params)
+	if err != nil {
+		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "解析导出表头失败")
+		h.logger.Error("解析导出表头失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
+		return err
+	}
+
+	shards, totalRows, err := h.buildShards(ctx, strategy, params, exportTask.CancelRequested, exportTask.ID, updater)
 	if err != nil {
 		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "构建导出分片失败")
 		h.logger.Error("构建导出分片失败", zap.Uint("task_id", exportTask.ID), zap.Error(err))
 		return err
 	}

-	if err := h.persistShardsAndPlan(ctx, exportTask.ID, updater, shards, totalRows); err != nil {
+	if err := h.persistShardsAndPlan(ctx, exportTask.ID, updater, shards, totalRows, headers); err != nil {
 		_ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "保存导出分片失败")
 		return err
 	}
@@ -141,58 +149,53 @@ func (h *ExportDispatchHandler) HandleExportDispatch(ctx context.Context, task *

 	h.logger.Info("导出 dispatch 完成",
 		zap.Uint("task_id", exportTask.ID),
-		zap.Uint("total_rows", totalRows),
+		zap.Int("total_rows", totalRows),
 		zap.Int("total_shards", len(shards)),
 	)

 	return nil
 }

-func (h *ExportDispatchHandler) buildShards(ctx context.Context, task *model.ExportTask, strategy exporter.SceneStrategy, updater uint) ([]*model.ExportShardTask, uint, error) {
+func (h *ExportDispatchHandler) buildShards(ctx context.Context, source exporter.DataSource, params exporter.ExportParams, cancelRequested bool, taskID uint, updater uint) ([]*model.ExportShardTask, int, error) {
 	shards := make([]*model.ExportShardTask, 0)
-	totalRows := uint(0)
-	afterID := uint64(0)
+	if cancelRequested {
+		return shards, 0, nil
+	}
+
+	totalRows, err := source.Count(ctx, params)
+	if err != nil {
+		return nil, 0, err
+	}
+	if totalRows <= 0 {
+		return shards, 0, nil
+	}
+
 	shardNo := 1
-
-	for {
-		if task.CancelRequested {
-			break
+	for offset := 0; offset < totalRows; offset += constants.ExportDefaultShardSize {
+		limit := constants.ExportDefaultShardSize
+		if remaining := totalRows - offset; remaining < limit {
+			limit = remaining
 		}

-		boundary, err := strategy.NextShardBoundary(ctx, task, afterID, constants.ExportDefaultShardSize)
-		if err != nil {
-			return nil, 0, err
-		}
-		if boundary == nil {
-			break
-		}
-
-		rowCount := boundary.RowCount
-		if rowCount < 0 {
-			rowCount = 0
-		}
-
-		totalRows += uint(rowCount)
 		shard := &model.ExportShardTask{
-			TaskID:      task.ID,
+			TaskID:      taskID,
 			ShardNo:     shardNo,
 			Status:      constants.ExportShardStatusPending,
-			CursorStart: boundary.StartID,
-			CursorEnd:   boundary.EndID,
-			RowCount:    rowCount,
+			ShardOffset: offset,
+			ShardLimit:  limit,
+			RowCount:    limit,
 		}
 		shard.Creator = updater
 		shard.Updater = updater
 		shards = append(shards, shard)

-		afterID = boundary.EndID
 		shardNo++
 	}

 	return shards, totalRows, nil
 }

-func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows uint) error {
+func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows int, headers []string) error {
 	return h.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
 		txShardStore := postgres.NewExportShardTaskStore(tx, h.redis)
 		txTaskStore := postgres.NewExportTaskStore(tx, h.redis)
@@ -200,7 +203,10 @@ func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID
 		if err := txShardStore.CreateBatch(ctx, shards); err != nil {
 			return err
 		}
-		if err := txTaskStore.UpdateDispatchPlan(ctx, taskID, totalRows, uint(len(shards)), updater); err != nil {
+		if err := txTaskStore.UpdateResolvedHeaders(ctx, taskID, headers, updater); err != nil {
+			return err
+		}
+		if err := txTaskStore.UpdateDispatchPlan(ctx, taskID, totalRows, len(shards), updater); err != nil {
 			return err
 		}
 		return nil
--- a/internal/task/export_finalize.go
+++ b/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())
--- a/internal/task/export_shard.go
+++ b/internal/task/export_shard.go
@@ -107,24 +107,38 @@ func (h *ExportShardHandler) HandleExportShard(ctx context.Context, task *asynq.
 		}
 	}

-	strategy, ok := h.sceneRegistry.Get(exportTask.Scene)
+	source, ok := h.sceneRegistry.Get(exportTask.Scene)
 	if !ok {
 		_ = h.markShardFinalFailed(ctx, exportTask.ID, shardTask.ID, updater, "导出场景未注册")
 		return asynq.SkipRetry
 	}

-	rows, err := strategy.QueryRows(ctx, exportTask, shardTask.CursorStart, shardTask.CursorEnd)
+	headers := exporter.ResolvedHeadersFromTask(exportTask)
+	if len(headers) == 0 {
+		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "读取导出表头失败", fmt.Errorf("resolved_headers 为空"))
+	}
+
+	params := exporter.ParseExportParams(exportTask)
+	rows, err := source.Fetch(ctx, params, shardTask.ShardOffset, shardTask.ShardLimit)
 	if err != nil {
 		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "查询分片数据失败", err)
 	}
+	if hasMisalignedRows(rows, headers) {
+		h.logger.Warn("导出分片数据列数与表头不一致，已按固定表头对齐",
+			zap.Uint("task_id", exportTask.ID),
+			zap.Uint("shard_id", shardTask.ID),
+			zap.Int("header_count", len(headers)),
+		)
+	}
+	rows = alignRowsToHeaders(rows, headers)

-	localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows)
+	localPath, fileSize, err := writeExportFile(constants.ExportTaskFormatCSV, headers, rows, true)
 	if err != nil {
 		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "生成分片文件失败", err)
 	}
 	defer func() { _ = os.Remove(localPath) }()

-	fileName := fmt.Sprintf("%s-shard-%04d.%s", exportTask.TaskNo, shardTask.ShardNo, exportTask.Format)
+	fileName := fmt.Sprintf("%s-shard-%04d.csv.shard", exportTask.TaskNo, shardTask.ShardNo)
 	fileKey, err := uploadExportFile(ctx, h.storageSvc, localPath, fileName)
 	if err != nil {
 		return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "上传分片文件失败", err)