统一导出任务

internal/task/export_common.go

@@ -0,0 +1,209 @@
+package task
+
+import (
+	"context"
+	"encoding/csv"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/bytedance/sonic"
+	"github.com/hibiken/asynq"
+	"github.com/redis/go-redis/v9"
+	"github.com/xuri/excelize/v2"
+
+	"github.com/break/junhong_cmp_fiber/pkg/constants"
+	"github.com/break/junhong_cmp_fiber/pkg/storage"
+)
+
+const (
+	exportLockTTL         = 5 * time.Minute
+	exportShardLockTTL    = 10 * time.Minute
+	exportFinalizeLockTTL = 2 * time.Minute
+)
+
+// ExportDispatchPayload 导出 dispatch 任务载荷。
+type ExportDispatchPayload struct {
+	TaskID uint `json:"task_id"`
+}
+
+// ExportShardPayload 导出 shard 任务载荷。
+type ExportShardPayload struct {
+	TaskID  uint `json:"task_id"`
+	ShardID uint `json:"shard_id"`
+}
+
+// ExportFinalizePayload 导出 finalize 任务载荷。
+type ExportFinalizePayload struct {
+	TaskID uint `json:"task_id"`
+}
+
+func enqueueTask(ctx context.Context, client *asynq.Client, taskType string, payload any, opts ...asynq.Option) error {
+	if client == nil {
+		return errors.New("asynq client 未初始化")
+	}
+
+	bytes, err := sonic.Marshal(payload)
+	if err != nil {
+		return err
+	}
+
+	task := asynq.NewTask(taskType, bytes, opts...)
+	_, err = client.EnqueueContext(ctx, task)
+	return err
+}
+
+func tryAcquireLock(ctx context.Context, rdb *redis.Client, key string, ttl time.Duration) (bool, error) {
+	if rdb == nil {
+		return true, nil
+	}
+	return rdb.SetNX(ctx, key, "1", ttl).Result()
+}
+
+func releaseLock(ctx context.Context, rdb *redis.Client, key string) {
+	if rdb == nil {
+		return
+	}
+	_, _ = rdb.Del(ctx, key).Result()
+}
+
+func isFinalRetry(ctx context.Context) bool {
+	retryCount, ok := asynq.GetRetryCount(ctx)
+	if !ok {
+		return false
+	}
+	maxRetry, ok := asynq.GetMaxRetry(ctx)
+	if !ok {
+		return false
+	}
+	return retryCount >= maxRetry
+}
+
+func writeExportFile(format string, headers []string, rows [][]string) (string, int64, error) {
+	switch format {
+	case constants.ExportTaskFormatCSV:
+		return writeCSVFile(headers, rows)
+	case constants.ExportTaskFormatXLSX:
+		return writeXLSXFile(headers, rows)
+	default:
+		return "", 0, fmt.Errorf("不支持的导出格式: %s", format)
+	}
+}
+
+func writeCSVFile(headers []string, rows [][]string) (string, int64, error) {
+	file, err := os.CreateTemp("", "export-*.csv")
+	if err != nil {
+		return "", 0, err
+	}
+	defer file.Close()
+
+	writer := csv.NewWriter(file)
+	if err := writer.Write(headers); err != nil {
+		return "", 0, err
+	}
+	for _, row := range rows {
+		if err := writer.Write(row); err != nil {
+			return "", 0, err
+		}
+	}
+	writer.Flush()
+	if err := writer.Error(); err != nil {
+		return "", 0, err
+	}
+
+	if err := file.Sync(); err != nil {
+		return "", 0, err
+	}
+
+	stat, err := file.Stat()
+	if err != nil {
+		return "", 0, err
+	}
+	return file.Name(), stat.Size(), nil
+}
+
+func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) {
+	file := excelize.NewFile()
+	sheet := file.GetSheetName(file.GetActiveSheetIndex())
+
+	streamWriter, err := file.NewStreamWriter(sheet)
+	if err != nil {
+		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
+	}
+
+	for i, row := range rows {
+		axis, err := excelize.CoordinatesToCellName(1, i+2)
+		if err != nil {
+			return "", 0, err
+		}
+		values := make([]interface{}, 0, len(row))
+		for _, item := range row {
+			values = append(values, item)
+		}
+		if err := streamWriter.SetRow(axis, values); err != nil {
+			return "", 0, err
+		}
+	}
+
+	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 uploadExportFile(ctx context.Context, storageSvc *storage.Service, localPath, fileName string) (string, error) {
+	if storageSvc == nil || storageSvc.Provider() == nil {
+		return "", errors.New("对象存储服务未配置")
+	}
+
+	key, err := storageSvc.GenerateFileKey("export", fileName)
+	if err != nil {
+		return "", err
+	}
+
+	file, err := os.Open(localPath)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	contentType := "application/octet-stream"
+	if filepath.Ext(fileName) == ".csv" {
+		contentType = "text/csv"
+	}
+	if filepath.Ext(fileName) == ".xlsx" {
+		contentType = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+	}
+
+	if err := storageSvc.Provider().Upload(ctx, key, file, contentType); err != nil {
+		return "", err
+	}
+	return key, nil
+}
+
+func parseExportPayload[T any](payload []byte, target *T) error {
+	return sonic.Unmarshal(payload, target)
+}