统一导出任务

2026-04-27 09:47:03 +08:00
parent 531de3c760
commit 60debb7505
29 changed files with 2580 additions and 0 deletions
--- a/internal/store/postgres/export_shard_task_store.go
+++ b/internal/store/postgres/export_shard_task_store.go
@@ -0,0 +1,154 @@
+package postgres
+
+import (
+	"context"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"gorm.io/gorm"
+
+	"github.com/break/junhong_cmp_fiber/internal/model"
+	"github.com/break/junhong_cmp_fiber/pkg/constants"
+)
+
+// ExportShardTaskStore 导出分片任务数据访问层。
+type ExportShardTaskStore struct {
+	db    *gorm.DB
+	redis *redis.Client
+}
+
+// NewExportShardTaskStore 创建导出分片任务 Store。
+func NewExportShardTaskStore(db *gorm.DB, redis *redis.Client) *ExportShardTaskStore {
+	return &ExportShardTaskStore{db: db, redis: redis}
+}
+
+// CreateBatch 批量创建分片任务。
+func (s *ExportShardTaskStore) CreateBatch(ctx context.Context, shards []*model.ExportShardTask) error {
+	if len(shards) == 0 {
+		return nil
+	}
+	return s.db.WithContext(ctx).Create(&shards).Error
+}
+
+// GetByIDForWorker 按 ID 查询分片任务（Worker 使用）。
+func (s *ExportShardTaskStore) GetByIDForWorker(ctx context.Context, id uint) (*model.ExportShardTask, error) {
+	var shard model.ExportShardTask
+	if err := s.db.WithContext(ctx).Where("id = ?", id).First(&shard).Error; err != nil {
+		return nil, err
+	}
+	return &shard, nil
+}
+
+// ListByTaskID 按任务 ID 查询分片任务列表。
+func (s *ExportShardTaskStore) ListByTaskID(ctx context.Context, taskID uint) ([]*model.ExportShardTask, error) {
+	var list []*model.ExportShardTask
+	if err := s.db.WithContext(ctx).
+		Where("task_id = ?", taskID).
+		Order("shard_no ASC").
+		Find(&list).Error; err != nil {
+		return nil, err
+	}
+	return list, nil
+}
+
+// TryMarkProcessingIfPending 尝试把分片从待处理更新为处理中。
+func (s *ExportShardTaskStore) TryMarkProcessingIfPending(ctx context.Context, shardID uint, updater uint) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("id = ? AND status = ?", shardID, constants.ExportShardStatusPending).
+		Updates(map[string]any{
+			"status":     constants.ExportShardStatusProcessing,
+			"started_at": now,
+			"updated_at": now,
+			"updater":    updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// MarkSuccess 标记分片为成功。
+func (s *ExportShardTaskStore) MarkSuccess(ctx context.Context, shardID uint, updater uint, rowCount int, fileKey string, fileSize int64) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
+		Updates(map[string]any{
+			"status":         constants.ExportShardStatusSuccess,
+			"row_count":      rowCount,
+			"processed_rows": rowCount,
+			"file_key":       fileKey,
+			"file_size":      fileSize,
+			"completed_at":   now,
+			"updated_at":     now,
+			"updater":        updater,
+			"error_message":  "",
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// MarkFailed 标记分片为失败。
+func (s *ExportShardTaskStore) MarkFailed(ctx context.Context, shardID uint, updater uint, errorMessage string) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
+		Updates(map[string]any{
+			"status":        constants.ExportShardStatusFailed,
+			"error_message": errorMessage,
+			"completed_at":  now,
+			"updated_at":    now,
+			"updater":       updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// MarkCancelled 标记分片为取消。
+func (s *ExportShardTaskStore) MarkCancelled(ctx context.Context, shardID uint, updater uint, errorMessage string) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
+		Updates(map[string]any{
+			"status":        constants.ExportShardStatusCancelled,
+			"error_message": errorMessage,
+			"completed_at":  now,
+			"updated_at":    now,
+			"updater":       updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// CancelByTaskID 批量取消指定主任务下未完成的分片。
+func (s *ExportShardTaskStore) CancelByTaskID(ctx context.Context, taskID uint, updater uint, errorMessage string) error {
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("task_id = ? AND status IN ?", taskID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
+		Updates(map[string]any{
+			"status":        constants.ExportShardStatusCancelled,
+			"error_message": errorMessage,
+			"completed_at":  now,
+			"updated_at":    now,
+			"updater":       updater,
+		}).Error
+}
+
+// CountUnfinishedByTaskID 统计未完成分片数量。
+func (s *ExportShardTaskStore) CountUnfinishedByTaskID(ctx context.Context, taskID uint) (int64, error) {
+	var count int64
+	err := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("task_id = ? AND status IN ?", taskID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
+		Count(&count).Error
+	return count, err
+}
+
+// CountByTaskIDAndStatus 统计指定状态分片数量。
+func (s *ExportShardTaskStore) CountByTaskIDAndStatus(ctx context.Context, taskID uint, status int) (int64, error) {
+	var count int64
+	err := s.db.WithContext(ctx).
+		Model(&model.ExportShardTask{}).
+		Where("task_id = ? AND status = ?", taskID, status).
+		Count(&count).Error
+	return count, err
+}
--- a/internal/store/postgres/export_task_store.go
+++ b/internal/store/postgres/export_task_store.go
@@ -0,0 +1,337 @@
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	"github.com/break/junhong_cmp_fiber/internal/model"
+	"github.com/break/junhong_cmp_fiber/internal/store"
+	"github.com/break/junhong_cmp_fiber/pkg/constants"
+	"github.com/break/junhong_cmp_fiber/pkg/middleware"
+)
+
+// ExportTaskStore 导出主任务数据访问层。
+type ExportTaskStore struct {
+	db    *gorm.DB
+	redis *redis.Client
+}
+
+// NewExportTaskStore 创建导出主任务 Store。
+func NewExportTaskStore(db *gorm.DB, redis *redis.Client) *ExportTaskStore {
+	return &ExportTaskStore{db: db, redis: redis}
+}
+
+// Create 创建导出主任务。
+func (s *ExportTaskStore) Create(ctx context.Context, task *model.ExportTask) error {
+	return s.db.WithContext(ctx).Create(task).Error
+}
+
+// GetByID 按 ID 查询任务（带数据权限）。
+func (s *ExportTaskStore) GetByID(ctx context.Context, id uint) (*model.ExportTask, error) {
+	var task model.ExportTask
+	query := applyExportTaskVisibility(ctx, s.db.WithContext(ctx).Model(&model.ExportTask{}))
+	if err := query.Where("id = ?", id).First(&task).Error; err != nil {
+		return nil, err
+	}
+	return &task, nil
+}
+
+// GetByIDForWorker 按 ID 查询任务（Worker 使用，不带数据权限）。
+func (s *ExportTaskStore) GetByIDForWorker(ctx context.Context, id uint) (*model.ExportTask, error) {
+	var task model.ExportTask
+	if err := s.db.WithContext(ctx).Where("id = ?", id).First(&task).Error; err != nil {
+		return nil, err
+	}
+	return &task, nil
+}
+
+// List 查询导出任务列表（带数据权限）。
+func (s *ExportTaskStore) List(ctx context.Context, opts *store.QueryOptions, filters map[string]any) ([]*model.ExportTask, int64, error) {
+	var (
+		items []*model.ExportTask
+		total int64
+	)
+
+	query := applyExportTaskVisibility(ctx, s.db.WithContext(ctx).Model(&model.ExportTask{}))
+
+	if scene, ok := filters["scene"].(string); ok && scene != "" {
+		query = query.Where("scene = ?", scene)
+	}
+	if status, ok := filters["status"].(int); ok && status > 0 {
+		query = query.Where("status = ?", status)
+	}
+	if startTime, ok := filters["start_time"].(time.Time); ok && !startTime.IsZero() {
+		query = query.Where("created_at >= ?", startTime)
+	}
+	if endTime, ok := filters["end_time"].(time.Time); ok && !endTime.IsZero() {
+		query = query.Where("created_at <= ?", endTime)
+	}
+
+	if err := query.Count(&total).Error; err != nil {
+		return nil, 0, err
+	}
+
+	if opts == nil {
+		opts = &store.QueryOptions{Page: 1, PageSize: constants.DefaultPageSize}
+	}
+	if opts.Page <= 0 {
+		opts.Page = 1
+	}
+	if opts.PageSize <= 0 {
+		opts.PageSize = constants.DefaultPageSize
+	}
+
+	offset := (opts.Page - 1) * opts.PageSize
+	query = query.Offset(offset).Limit(opts.PageSize)
+
+	if opts.OrderBy != "" {
+		query = query.Order(opts.OrderBy)
+	} else {
+		query = query.Order("created_at DESC")
+	}
+
+	if err := query.Find(&items).Error; err != nil {
+		return nil, 0, err
+	}
+
+	return items, total, nil
+}
+
+// GenerateTaskNo 生成导出任务编号。
+func (s *ExportTaskStore) GenerateTaskNo() string {
+	now := time.Now()
+	return fmt.Sprintf("EXP-%s-%06d", now.Format("20060102"), now.UnixNano()%1000000)
+}
+
+// TryMarkProcessingIfPending 尝试将任务从待处理更新为处理中。
+func (s *ExportTaskStore) TryMarkProcessingIfPending(ctx context.Context, taskID uint, updater uint) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status = ?", taskID, constants.ExportTaskStatusPending).
+		Updates(map[string]any{
+			"status":     constants.ExportTaskStatusProcessing,
+			"started_at": now,
+			"updated_at": now,
+			"updater":    updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// UpdateDispatchPlan 更新 dispatch 阶段计算出的总量信息。
+func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, totalRows, totalShards uint, updater uint) error {
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ?", taskID).
+		Updates(map[string]any{
+			"total_rows":    totalRows,
+			"total_shards":  totalShards,
+			"updated_at":    now,
+			"updater":       updater,
+			"error_message": "",
+		}).Error
+}
+
+// IsCancelRequested 查询是否已请求取消。
+func (s *ExportTaskStore) IsCancelRequested(ctx context.Context, taskID uint) (bool, error) {
+	var value bool
+	if err := s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ?", taskID).
+		Pluck("cancel_requested", &value).Error; err != nil {
+		return false, err
+	}
+	return value, nil
+}
+
+// CancelPendingTask 将待处理任务直接取消。
+func (s *ExportTaskStore) CancelPendingTask(ctx context.Context, taskID uint, updater uint) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status = ?", taskID, constants.ExportTaskStatusPending).
+		Updates(map[string]any{
+			"status":           constants.ExportTaskStatusCancelled,
+			"cancel_requested": true,
+			"progress":         100,
+			"completed_at":     now,
+			"updated_at":       now,
+			"updater":          updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// SetCancelRequested 标记处理中任务取消请求。
+func (s *ExportTaskStore) SetCancelRequested(ctx context.Context, taskID uint, updater uint) (bool, error) {
+	now := time.Now()
+	result := s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status = ?", taskID, constants.ExportTaskStatusProcessing).
+		Updates(map[string]any{
+			"cancel_requested": true,
+			"updated_at":       now,
+			"updater":          updater,
+		})
+	return result.RowsAffected > 0, result.Error
+}
+
+// MarkCancelledByWorker 在 Worker 侧将任务状态更新为已取消。
+func (s *ExportTaskStore) MarkCancelledByWorker(ctx context.Context, taskID uint, updater uint, msg string) error {
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status IN ?", taskID, []int{constants.ExportTaskStatusPending, constants.ExportTaskStatusProcessing}).
+		Updates(map[string]any{
+			"status":           constants.ExportTaskStatusCancelled,
+			"cancel_requested": true,
+			"progress":         100,
+			"completed_at":     now,
+			"updated_at":       now,
+			"updater":          updater,
+			"error_message":    msg,
+		}).Error
+}
+
+// MarkCompleted 将任务更新为已完成。
+func (s *ExportTaskStore) MarkCompleted(ctx context.Context, taskID uint, updater uint, fileKey string, fileSize int64) error {
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status = ?", taskID, constants.ExportTaskStatusProcessing).
+		Updates(map[string]any{
+			"status":        constants.ExportTaskStatusCompleted,
+			"progress":      100,
+			"file_key":      fileKey,
+			"file_size":     fileSize,
+			"completed_at":  now,
+			"updated_at":    now,
+			"updater":       updater,
+			"error_message": "",
+		}).Error
+}
+
+// MarkFailed 将任务更新为失败。
+func (s *ExportTaskStore) MarkFailed(ctx context.Context, taskID uint, updater uint, errorMessage string) error {
+	now := time.Now()
+	return s.db.WithContext(ctx).
+		Model(&model.ExportTask{}).
+		Where("id = ? AND status IN ?", taskID, []int{constants.ExportTaskStatusPending, constants.ExportTaskStatusProcessing}).
+		Updates(map[string]any{
+			"status":        constants.ExportTaskStatusFailed,
+			"completed_at":  now,
+			"updated_at":    now,
+			"updater":       updater,
+			"error_message": errorMessage,
+		}).Error
+}
+
+// ApplyShardSuccess 在分片成功后更新主任务进度与统计。
+func (s *ExportTaskStore) ApplyShardSuccess(ctx context.Context, taskID uint, updater uint, rowCount int) error {
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var task model.ExportTask
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
+			Where("id = ?", taskID).
+			First(&task).Error; err != nil {
+			return err
+		}
+		if task.Status != constants.ExportTaskStatusProcessing {
+			return nil
+		}
+
+		processedRows := task.ProcessedRows + rowCount
+		successShards := task.SuccessShards + 1
+		progress := task.Progress
+
+		if task.TotalRows > 0 {
+			progress = processedRows * 100 / task.TotalRows
+		} else if task.TotalShards > 0 {
+			progress = successShards * 100 / task.TotalShards
+		}
+		if progress > 99 {
+			progress = 99
+		}
+		if progress < task.Progress {
+			progress = task.Progress
+		}
+
+		now := time.Now()
+		return tx.Model(&model.ExportTask{}).
+			Where("id = ?", taskID).
+			Updates(map[string]any{
+				"processed_rows": processedRows,
+				"success_shards": successShards,
+				"progress":       progress,
+				"updated_at":     now,
+				"updater":        updater,
+			}).Error
+	})
+}
+
+// ApplyShardFailed 在分片失败后更新主任务失败计数。
+func (s *ExportTaskStore) ApplyShardFailed(ctx context.Context, taskID uint, updater uint, errorMessage string) error {
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		var task model.ExportTask
+		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
+			Where("id = ?", taskID).
+			First(&task).Error; err != nil {
+			return err
+		}
+		if task.Status != constants.ExportTaskStatusProcessing {
+			return nil
+		}
+
+		now := time.Now()
+		updates := map[string]any{
+			"failed_shards": task.FailedShards + 1,
+			"updated_at":    now,
+			"updater":       updater,
+		}
+		if task.ErrorMessage == "" {
+			updates["error_message"] = errorMessage
+		}
+
+		return tx.Model(&model.ExportTask{}).
+			Where("id = ?", taskID).
+			Updates(updates).Error
+	})
+}
+
+func applyExportTaskVisibility(ctx context.Context, query *gorm.DB) *gorm.DB {
+	userType := middleware.GetUserTypeFromContext(ctx)
+
+	switch userType {
+	case constants.UserTypeSuperAdmin, constants.UserTypePlatform:
+		return query
+	case constants.UserTypeAgent:
+		shopIDs := middleware.GetSubordinateShopIDs(ctx)
+		if shopIDs == nil {
+			shopID := middleware.GetShopIDFromContext(ctx)
+			if shopID == 0 {
+				return query.Where("1 = 0")
+			}
+			return query.Where("creator_shop_id = ?", shopID)
+		}
+		if len(shopIDs) == 0 {
+			return query.Where("1 = 0")
+		}
+		return query.Where("creator_shop_id IN ?", shopIDs)
+	case constants.UserTypeEnterprise:
+		enterpriseID := middleware.GetEnterpriseIDFromContext(ctx)
+		if enterpriseID == 0 {
+			return query.Where("1 = 0")
+		}
+		return query.Where("creator_enterprise_id = ?", enterpriseID)
+	default:
+		userID := middleware.GetUserIDFromContext(ctx)
+		if userID == 0 {
+			return query.Where("1 = 0")
+		}
+		return query.Where("creator_user_id = ?", userID)
+	}
+}