7ec84fbc0f
All checks were successful
构建并部署到测试环境（无 SSH） / build-and-deploy (push) Successful in 7m50s
259 lines
6.9 KiB
Go
259 lines
6.9 KiB
Go
package task
|
|
|
|
import (
|
|
"context"
|
|
"encoding/csv"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
|
|
"github.com/hibiken/asynq"
|
|
"github.com/redis/go-redis/v9"
|
|
"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"
|
|
)
|
|
|
|
// ExportFinalizeHandler 导出 finalize 处理器。
|
|
type ExportFinalizeHandler struct {
|
|
redis *redis.Client
|
|
taskStore *postgres.ExportTaskStore
|
|
shardStore *postgres.ExportShardTaskStore
|
|
storageSvc *storage.Service
|
|
sceneRegistry *exporter.Registry
|
|
logger *zap.Logger
|
|
}
|
|
|
|
// NewExportFinalizeHandler 创建导出 finalize 处理器。
|
|
func NewExportFinalizeHandler(
|
|
redis *redis.Client,
|
|
taskStore *postgres.ExportTaskStore,
|
|
shardStore *postgres.ExportShardTaskStore,
|
|
storageSvc *storage.Service,
|
|
sceneRegistry *exporter.Registry,
|
|
logger *zap.Logger,
|
|
) *ExportFinalizeHandler {
|
|
return &ExportFinalizeHandler{
|
|
redis: redis,
|
|
taskStore: taskStore,
|
|
shardStore: shardStore,
|
|
storageSvc: storageSvc,
|
|
sceneRegistry: sceneRegistry,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
// HandleExportFinalize 处理导出 finalize 任务。
|
|
func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task *asynq.Task) error {
|
|
var payload ExportFinalizePayload
|
|
if err := parseExportPayload(task.Payload(), &payload); err != nil {
|
|
h.logger.Error("解析导出 finalize 任务载荷失败", zap.Error(err))
|
|
return asynq.SkipRetry
|
|
}
|
|
|
|
lockKey := constants.RedisExportFinalizeLockKey(payload.TaskID)
|
|
locked, err := tryAcquireLock(ctx, h.redis, lockKey, exportFinalizeLockTTL)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !locked {
|
|
return nil
|
|
}
|
|
defer releaseLock(ctx, h.redis, lockKey)
|
|
|
|
exportTask, err := h.taskStore.GetByIDForWorker(ctx, payload.TaskID)
|
|
if err != nil {
|
|
h.logger.Error("查询导出任务失败", zap.Uint("task_id", payload.TaskID), zap.Error(err))
|
|
return asynq.SkipRetry
|
|
}
|
|
|
|
if exportTask.Status == constants.ExportTaskStatusCompleted || exportTask.Status == constants.ExportTaskStatusFailed || exportTask.Status == constants.ExportTaskStatusCancelled {
|
|
return nil
|
|
}
|
|
|
|
updater := dispatchUpdater(exportTask)
|
|
if exportTask.CancelRequested {
|
|
if err := h.shardStore.CancelByTaskID(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
|
|
return err
|
|
}
|
|
if err := h.taskStore.MarkCancelledByWorker(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
unfinishedCount, err := h.shardStore.CountUnfinishedByTaskID(ctx, exportTask.ID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if unfinishedCount > 0 {
|
|
return fmt.Errorf("分片任务尚未完成: %d", unfinishedCount)
|
|
}
|
|
|
|
failedCount, err := h.shardStore.CountByTaskIDAndStatus(ctx, exportTask.ID, constants.ExportShardStatusFailed)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if failedCount > 0 {
|
|
if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "存在失败分片，导出任务失败"); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
if !h.sceneRegistry.IsSupported(exportTask.Scene) {
|
|
if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "导出场景未注册"); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
shards, err := h.shardStore.ListByTaskID(ctx, exportTask.ID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
headers := exporter.ResolvedHeadersFromTask(exportTask)
|
|
if len(headers) == 0 {
|
|
return h.handleFinalizeError(ctx, exportTask.ID, updater, "读取导出表头失败", errors.New("resolved_headers 为空"))
|
|
}
|
|
|
|
mergedCSVPath, cleanupCSV, err := h.mergeShardCSVFiles(ctx, headers, shards)
|
|
if err != nil {
|
|
return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
|
|
}
|
|
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)
|
|
if err != nil {
|
|
return h.handleFinalizeError(ctx, exportTask.ID, updater, "上传汇总文件失败", err)
|
|
}
|
|
|
|
if err := h.taskStore.MarkCompleted(ctx, exportTask.ID, updater, fileKey, fileSize); err != nil {
|
|
return err
|
|
}
|
|
|
|
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())
|
|
if markErr != nil {
|
|
return markErr
|
|
}
|
|
return nil
|
|
}
|
|
return err
|
|
}
|