junhong_cmp_fiber/internal/task/export_common.go

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
	}

	opts = append(opts, asynq.Queue(constants.QueueForTaskType(taskType)))
	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)
}