统一导出任务

internal/task/export_finalize.go

@@ -0,0 +1,158 @@
+package task
+
+import (
+	"context"
+	"fmt"
+	"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/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
+	}
+
+	strategy, ok := h.sceneRegistry.Get(exportTask.Scene)
+	if !ok {
+		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
+	}
+
+	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...)
+	}
+
+	localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows)
+	if err != nil {
+		return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
+	}
+	defer func() { _ = os.Remove(localPath) }()
+
+	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) 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
+}