From 7ec84fbc0f7f5009a98c3660a9521504f8d0f4a4 Mon Sep 17 00:00:00 2001 From: Break Date: Mon, 15 Jun 2026 15:28:29 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AF=BC=E5=87=BA=E7=B3=BB=E7=BB=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/exporter/datasource.go | 20 ++ internal/exporter/device_scene.go | 100 +++---- internal/exporter/filter_helpers.go | 270 ++++++++++++++++++ internal/exporter/iot_card_scene.go | 248 ++++++++++++---- internal/exporter/query_params.go | 56 ++++ internal/exporter/registry.go | 55 ++-- internal/exporter/scope.go | 25 -- internal/model/export_task.go | 10 +- internal/store/postgres/export_task_store.go | 38 ++- internal/task/export_common.go | 122 +++++++- internal/task/export_dispatch.go | 68 +++-- internal/task/export_finalize.go | 128 ++++++++- internal/task/export_shard.go | 22 +- ...53_alter_export_shard_task_offset.down.sql | 4 + ...0153_alter_export_shard_task_offset.up.sql | 7 + pkg/storage/s3.go | 24 +- pkg/storage/storage.go | 2 +- 17 files changed, 939 insertions(+), 260 deletions(-) create mode 100644 internal/exporter/datasource.go create mode 100644 internal/exporter/filter_helpers.go create mode 100644 internal/exporter/query_params.go create mode 100644 migrations/000153_alter_export_shard_task_offset.down.sql create mode 100644 migrations/000153_alter_export_shard_task_offset.up.sql diff --git a/internal/exporter/datasource.go b/internal/exporter/datasource.go new file mode 100644 index 0000000..5d5d1c8 --- /dev/null +++ b/internal/exporter/datasource.go @@ -0,0 +1,20 @@ +package exporter + +import "context" + +// DataSource 导出场景数据源接口。 +// 场景实现只负责统计、表头和按 offset/limit 拉取数据,分片调度由导出框架统一处理。 +type DataSource interface { + Scene() string + Count(ctx context.Context, params ExportParams) (int, error) + Headers(ctx context.Context, params ExportParams) ([]string, error) + Fetch(ctx context.Context, params ExportParams, offset, limit int) ([][]string, error) +} + +// ExportParams 导出任务执行参数。 +// Filters 来自任务 query_json.filters,权限字段来自任务创建时的快照。 +type ExportParams struct { + Filters map[string]any + ScopeShopIDs []uint + UserType int +} diff --git a/internal/exporter/device_scene.go b/internal/exporter/device_scene.go index 06ed3bb..af04798 100644 --- a/internal/exporter/device_scene.go +++ b/internal/exporter/device_scene.go @@ -10,78 +10,53 @@ import ( "github.com/break/junhong_cmp_fiber/pkg/constants" ) -// DeviceSceneStrategy 设备导出场景策略。 -type DeviceSceneStrategy struct { +// DeviceDataSource 设备导出数据源。 +type DeviceDataSource struct { db *gorm.DB } -// NewDeviceSceneStrategy 创建设备场景策略。 -func NewDeviceSceneStrategy(db *gorm.DB) *DeviceSceneStrategy { - return &DeviceSceneStrategy{db: db} +// NewDeviceDataSource 创建设备导出数据源。 +func NewDeviceDataSource(db *gorm.DB) *DeviceDataSource { + return &DeviceDataSource{db: db} } -// Scene 场景编码。 -func (s *DeviceSceneStrategy) Scene() string { +// Scene 返回导出场景编码。 +func (s *DeviceDataSource) Scene() string { return constants.ExportTaskSceneDevice } -// Headers 导出表头。 -func (s *DeviceSceneStrategy) Headers() []string { - return []string{"ID", "设备虚拟号", "设备名称", "设备型号", "设备类型", "状态", "店铺ID", "创建时间"} +// Count 统计设备导出行数。 +func (s *DeviceDataSource) Count(ctx context.Context, params ExportParams) (int, error) { + var total int64 + query := s.applyFilters(s.db.WithContext(ctx).Model(&model.Device{}), params) + if err := query.Count(&total).Error; err != nil { + return 0, err + } + return int(total), nil } -// NextShardBoundary 获取下一段分片边界。 -func (s *DeviceSceneStrategy) NextShardBoundary(ctx context.Context, task *model.ExportTask, afterID uint64, limit int) (*ShardBoundary, error) { +// Headers 返回设备导出表头。 +func (s *DeviceDataSource) Headers(ctx context.Context, params ExportParams) ([]string, error) { + return []string{"ID", "设备虚拟号", "设备名称", "设备型号", "设备类型", "状态", "店铺ID", "创建时间"}, nil +} + +// Fetch 按 offset/limit 查询设备导出数据。 +func (s *DeviceDataSource) Fetch(ctx context.Context, params ExportParams, offset, limit int) ([][]string, error) { if limit <= 0 { - limit = constants.ExportDefaultShardSize - } - - var ids []uint64 - query := s.db.WithContext(ctx). - Model(&model.Device{}). - Where("id > ?", afterID). - Order("id ASC"). - Limit(limit) - query = applyShopScopeForTask(query, task) - - if err := query.Pluck("id", &ids).Error; err != nil { - return nil, err - } - if len(ids) == 0 { - return nil, nil - } - - return &ShardBoundary{ - StartID: afterID, - EndID: ids[len(ids)-1], - RowCount: len(ids), - }, nil -} - -// QueryRows 查询分片内数据行。 -func (s *DeviceSceneStrategy) QueryRows(ctx context.Context, task *model.ExportTask, startID, endID uint64) ([][]string, error) { - if endID == 0 { return [][]string{}, nil } var devices []model.Device - query := s.db.WithContext(ctx). - Model(&model.Device{}). - Where("id > ? AND id <= ?", startID, endID). - Order("id ASC") - query = applyShopScopeForTask(query, task) - + query := s.applyFilters(s.db.WithContext(ctx).Model(&model.Device{}), params). + Order("id ASC"). + Limit(limit). + Offset(offset) if err := query.Find(&devices).Error; err != nil { return nil, err } rows := make([][]string, 0, len(devices)) for _, item := range devices { - shopID := "" - if item.ShopID != nil { - shopID = strconv.FormatUint(uint64(*item.ShopID), 10) - } - rows = append(rows, []string{ strconv.FormatUint(uint64(item.ID), 10), item.VirtualNo, @@ -89,10 +64,27 @@ func (s *DeviceSceneStrategy) QueryRows(ctx context.Context, task *model.ExportT item.DeviceModel, item.DeviceType, strconv.Itoa(item.Status), - shopID, - item.CreatedAt.Format("2006-01-02 15:04:05"), + formatOptionalUint(item.ShopID), + item.CreatedAt.Format(exportTimeLayout), }) } - return rows, nil } + +func (s *DeviceDataSource) applyFilters(query *gorm.DB, params ExportParams) *gorm.DB { + query = applyExportShopScope(query, params, "shop_id") + + if status, ok := filterInt(params.Filters, "status"); ok { + query = query.Where("status = ?", status) + } + if shopID, ok := filterUint(params.Filters, "shop_id"); ok { + query = query.Where("shop_id = ?", shopID) + } + if start, ok := filterTime(params.Filters, "created_at_start"); ok { + query = query.Where("created_at >= ?", start) + } + if end, ok := filterTime(params.Filters, "created_at_end"); ok { + query = query.Where("created_at <= ?", end) + } + return query +} diff --git a/internal/exporter/filter_helpers.go b/internal/exporter/filter_helpers.go new file mode 100644 index 0000000..e80c2cf --- /dev/null +++ b/internal/exporter/filter_helpers.go @@ -0,0 +1,270 @@ +package exporter + +import ( + "fmt" + "strconv" + "strings" + "time" + + "gorm.io/gorm" + + "github.com/break/junhong_cmp_fiber/pkg/constants" +) + +const exportTimeLayout = "2006-01-02 15:04:05" + +func applyExportShopScope(query *gorm.DB, params ExportParams, column string) *gorm.DB { + switch params.UserType { + case constants.UserTypeSuperAdmin, constants.UserTypePlatform: + return query + case constants.UserTypeAgent: + if len(params.ScopeShopIDs) == 0 { + return query.Where("1 = 0") + } + return query.Where(column+" IN ?", params.ScopeShopIDs) + default: + return query.Where("1 = 0") + } +} + +func filterInt(filters map[string]any, key string) (int, bool) { + value, ok := filters[key] + if !ok || value == nil { + return 0, false + } + + switch v := value.(type) { + case int: + return v, true + case int64: + return int(v), true + case uint: + return int(v), true + case uint64: + return int(v), true + case float64: + return int(v), true + case string: + if strings.TrimSpace(v) == "" { + return 0, false + } + parsed, err := strconv.Atoi(v) + if err != nil { + return 0, false + } + return parsed, true + default: + return 0, false + } +} + +func filterUint(filters map[string]any, key string) (uint, bool) { + value, ok := filterInt(filters, key) + if !ok || value <= 0 { + return 0, false + } + return uint(value), true +} + +func filterString(filters map[string]any, key string) (string, bool) { + value, ok := filters[key] + if !ok || value == nil { + return "", false + } + text := strings.TrimSpace(fmt.Sprint(value)) + if text == "" { + return "", false + } + return text, true +} + +func filterBool(filters map[string]any, key string) bool { + value, ok := filters[key] + if !ok || value == nil { + return false + } + + switch v := value.(type) { + case bool: + return v + case string: + parsed, err := strconv.ParseBool(strings.TrimSpace(v)) + return err == nil && parsed + default: + return false + } +} + +// filterOptionalBool 解析可选布尔筛选项。 +// 返回第二个值标识调用方是否明确传入该筛选项,避免 false 被误判为未传。 +func filterOptionalBool(filters map[string]any, key string) (bool, bool) { + value, ok := filters[key] + if !ok || value == nil { + return false, false + } + + switch v := value.(type) { + case bool: + return v, true + case string: + if strings.TrimSpace(v) == "" { + return false, false + } + parsed, err := strconv.ParseBool(strings.TrimSpace(v)) + if err != nil { + return false, false + } + return parsed, true + case int: + return v != 0, true + case int64: + return v != 0, true + case uint: + return v != 0, true + case uint64: + return v != 0, true + case float64: + return v != 0, true + default: + return false, false + } +} + +// filterUintSlice 解析前端 JSON 或查询参数中的 uint 数组筛选项。 +// 空数组按未传处理;数组存在但过滤后为空时由调用方决定是否返回空结果。 +func filterUintSlice(filters map[string]any, key string) ([]uint, bool) { + value, ok := filters[key] + if !ok || value == nil { + return nil, false + } + + switch v := value.(type) { + case []uint: + if len(v) == 0 { + return nil, false + } + return normalizeUintSlice(v), true + case []int: + if len(v) == 0 { + return nil, false + } + result := make([]uint, 0, len(v)) + for _, item := range v { + if item > 0 { + result = append(result, uint(item)) + } + } + return normalizeUintSlice(result), true + case []float64: + if len(v) == 0 { + return nil, false + } + result := make([]uint, 0, len(v)) + for _, item := range v { + if item > 0 { + result = append(result, uint(item)) + } + } + return normalizeUintSlice(result), true + case []any: + if len(v) == 0 { + return nil, false + } + result := make([]uint, 0, len(v)) + for _, item := range v { + switch typed := item.(type) { + case int: + if typed > 0 { + result = append(result, uint(typed)) + } + case int64: + if typed > 0 { + result = append(result, uint(typed)) + } + case uint: + if typed > 0 { + result = append(result, typed) + } + case uint64: + if typed > 0 { + result = append(result, uint(typed)) + } + case float64: + if typed > 0 { + result = append(result, uint(typed)) + } + case string: + if parsed, err := strconv.ParseUint(strings.TrimSpace(typed), 10, 64); err == nil && parsed > 0 { + result = append(result, uint(parsed)) + } + } + } + return normalizeUintSlice(result), true + case string: + if strings.TrimSpace(v) == "" { + return nil, false + } + parts := strings.Split(v, ",") + result := make([]uint, 0, len(parts)) + for _, part := range parts { + parsed, err := strconv.ParseUint(strings.TrimSpace(part), 10, 64) + if err == nil && parsed > 0 { + result = append(result, uint(parsed)) + } + } + return normalizeUintSlice(result), true + default: + return nil, false + } +} + +// normalizeUintSlice 去除无效值和重复 ID,保留原始顺序。 +func normalizeUintSlice(values []uint) []uint { + seen := make(map[uint]struct{}, len(values)) + result := make([]uint, 0, len(values)) + for _, value := range values { + if value == 0 { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + result = append(result, value) + } + return result +} + +func filterTime(filters map[string]any, key string) (time.Time, bool) { + text, ok := filterString(filters, key) + if !ok { + return time.Time{}, false + } + + layouts := []string{ + time.RFC3339, + "2006-01-02 15:04:05", + "2006-01-02", + } + for _, layout := range layouts { + parsed, err := time.ParseInLocation(layout, text, time.Local) + if err == nil { + return parsed, true + } + } + return time.Time{}, false +} + +func formatOptionalUint(value *uint) string { + if value == nil { + return "" + } + return strconv.FormatUint(uint64(*value), 10) +} + +func formatOptionalTime(value *time.Time) string { + if value == nil { + return "" + } + return value.Format(exportTimeLayout) +} diff --git a/internal/exporter/iot_card_scene.go b/internal/exporter/iot_card_scene.go index a709177..644677b 100644 --- a/internal/exporter/iot_card_scene.go +++ b/internal/exporter/iot_card_scene.go @@ -2,97 +2,219 @@ package exporter import ( "context" - "strconv" + "time" "gorm.io/gorm" - "github.com/break/junhong_cmp_fiber/internal/model" "github.com/break/junhong_cmp_fiber/pkg/constants" ) -// IotCardSceneStrategy IoT 卡导出场景策略。 -type IotCardSceneStrategy struct { +// IotCardDataSource IoT 卡导出数据源。 +type IotCardDataSource struct { db *gorm.DB } -// NewIotCardSceneStrategy 创建 IoT 卡场景策略。 -func NewIotCardSceneStrategy(db *gorm.DB) *IotCardSceneStrategy { - return &IotCardSceneStrategy{db: db} +// NewIotCardDataSource 创建 IoT 卡导出数据源。 +func NewIotCardDataSource(db *gorm.DB) *IotCardDataSource { + return &IotCardDataSource{db: db} } -// Scene 场景编码。 -func (s *IotCardSceneStrategy) Scene() string { +// Scene 返回导出场景编码。 +func (s *IotCardDataSource) Scene() string { return constants.ExportTaskSceneIotCard } -// Headers 导出表头。 -func (s *IotCardSceneStrategy) Headers() []string { - return []string{"ID", "ICCID", "MSISDN", "虚拟号", "运营商名称", "状态", "店铺ID", "创建时间"} +// Count 统计 IoT 卡导出行数。 +func (s *IotCardDataSource) Count(ctx context.Context, params ExportParams) (int, error) { + var total int64 + query := s.applyFilters(ctx, s.baseQuery(ctx), params) + if err := query.Count(&total).Error; err != nil { + return 0, err + } + return int(total), nil } -// NextShardBoundary 获取下一段分片边界。 -func (s *IotCardSceneStrategy) NextShardBoundary(ctx context.Context, task *model.ExportTask, afterID uint64, limit int) (*ShardBoundary, error) { +// Headers 返回 IoT 卡导出表头。 +func (s *IotCardDataSource) Headers(ctx context.Context, params ExportParams) ([]string, error) { + return []string{"ICCID", "MSISDN", "绑定设备虚拟号", "运营商", "店铺名称", "绑定设备名称", "是否实名", "实名时间", "网络状态"}, nil +} + +// Fetch 按 offset/limit 查询 IoT 卡导出数据。 +func (s *IotCardDataSource) Fetch(ctx context.Context, params ExportParams, offset, limit int) ([][]string, error) { if limit <= 0 { - limit = constants.ExportDefaultShardSize - } - - var ids []uint64 - query := s.db.WithContext(ctx). - Model(&model.IotCard{}). - Where("id > ?", afterID). - Order("id ASC"). - Limit(limit) - query = applyShopScopeForTask(query, task) - - if err := query.Pluck("id", &ids).Error; err != nil { - return nil, err - } - if len(ids) == 0 { - return nil, nil - } - - return &ShardBoundary{ - StartID: afterID, - EndID: ids[len(ids)-1], - RowCount: len(ids), - }, nil -} - -// QueryRows 查询分片内数据行。 -func (s *IotCardSceneStrategy) QueryRows(ctx context.Context, task *model.ExportTask, startID, endID uint64) ([][]string, error) { - if endID == 0 { return [][]string{}, nil } - var cards []model.IotCard - query := s.db.WithContext(ctx). - Model(&model.IotCard{}). - Where("id > ? AND id <= ?", startID, endID). - Order("id ASC") - query = applyShopScopeForTask(query, task) - - if err := query.Find(&cards).Error; err != nil { + var items []iotCardExportRow + query := s.applyFilters(ctx, s.baseQuery(ctx), params). + Select(` + c.iccid, + c.msisdn, + c.device_virtual_no, + c.carrier_name, + COALESCE(sh.shop_name, '') AS shop_name, + COALESCE(d.device_name, '') AS device_name, + c.real_name_status, + c.first_realname_at, + c.network_status + `). + Order("c.id ASC"). + Limit(limit). + Offset(offset) + if err := query.Scan(&items).Error; err != nil { return nil, err } - rows := make([][]string, 0, len(cards)) - for _, item := range cards { - shopID := "" - if item.ShopID != nil { - shopID = strconv.FormatUint(uint64(*item.ShopID), 10) - } - + rows := make([][]string, 0, len(items)) + for _, item := range items { rows = append(rows, []string{ - strconv.FormatUint(uint64(item.ID), 10), item.ICCID, item.MSISDN, - item.VirtualNo, + item.DeviceVirtualNo, item.CarrierName, - strconv.Itoa(item.Status), - shopID, - item.CreatedAt.Format("2006-01-02 15:04:05"), + item.ShopName, + item.DeviceName, + formatRealNameVerified(item.RealNameStatus), + formatOptionalTime(item.FirstRealnameAt), + constants.GetNetworkStatusName(item.NetworkStatus), }) } - return rows, nil } + +// baseQuery 构造卡导出的基础查询。 +// 绑定设备名称从当前有效绑定中取一条,避免多绑定历史记录放大导出行数。 +func (s *IotCardDataSource) baseQuery(ctx context.Context) *gorm.DB { + return s.db.WithContext(ctx). + Table("tb_iot_card AS c"). + Joins("LEFT JOIN tb_shop AS sh ON sh.id = c.shop_id"). + Joins(` + LEFT JOIN LATERAL ( + SELECT d.device_name + FROM tb_device_sim_binding AS b + JOIN tb_device AS d ON d.id = b.device_id AND d.deleted_at IS NULL + WHERE b.iot_card_id = c.id + AND b.bind_status = ? + AND b.deleted_at IS NULL + ORDER BY b.is_current DESC, b.id DESC + LIMIT 1 + ) AS d ON TRUE + `, constants.BindStatusBound). + Where("c.deleted_at IS NULL") +} + +// applyFilters 应用与 /api/admin/iot-cards/standalone 一致的筛选条件。 +// 导出任务运行在 Worker 中,权限只能使用任务创建时保存的店铺范围快照。 +func (s *IotCardDataSource) applyFilters(ctx context.Context, query *gorm.DB, params ExportParams) *gorm.DB { + query = applyExportShopScope(query, params, "c.shop_id") + + if isStandalone, ok := filterOptionalBool(params.Filters, "is_standalone"); ok { + query = query.Where("c.is_standalone = ?", isStandalone) + } + if status, ok := filterInt(params.Filters, "status"); ok && status > 0 { + query = query.Where("c.status = ?", status) + } + if carrierID, ok := filterUint(params.Filters, "carrier_id"); ok { + query = query.Where("c.carrier_id = ?", carrierID) + } + if shopIDs, ok := filterUintSlice(params.Filters, "shop_ids"); ok { + if len(shopIDs) == 0 { + query = query.Where("1 = 0") + } else { + query = query.Where("c.shop_id IN ?", shopIDs) + } + } else if shopID, ok := filterInt(params.Filters, "shop_id"); ok { + if shopID <= 0 { + query = query.Where("1 = 0") + } else { + query = query.Where("c.shop_id = ?", shopID) + } + } + if iccid, ok := filterString(params.Filters, "iccid"); ok { + query = query.Where("c.iccid LIKE ?", "%"+iccid+"%") + } + if msisdn, ok := filterString(params.Filters, "msisdn"); ok { + query = query.Where("c.msisdn LIKE ?", "%"+msisdn+"%") + } + if virtualNo, ok := filterString(params.Filters, "virtual_no"); ok { + query = query.Where("c.virtual_no LIKE ?", "%"+virtualNo+"%") + } + if keyword, ok := filterString(params.Filters, "keyword"); ok { + query = query.Where("c.iccid LIKE ? OR c.virtual_no LIKE ?", "%"+keyword+"%", "%"+keyword+"%") + } + if batchNo, ok := filterString(params.Filters, "batch_no"); ok { + query = query.Where("c.batch_no = ?", batchNo) + } + if packageID, ok := filterUint(params.Filters, "package_id"); ok { + query = query.Where("c.id IN (?)", + s.db.WithContext(ctx).Table("tb_package_usage"). + Select("iot_card_id"). + Where("package_id = ? AND deleted_at IS NULL", packageID)) + } + if isDistributed, ok := filterOptionalBool(params.Filters, "is_distributed"); ok { + if isDistributed { + query = query.Where("c.shop_id IS NOT NULL") + } else { + query = query.Where("c.shop_id IS NULL") + } + } + if iccidStart, ok := filterString(params.Filters, "iccid_start"); ok { + query = query.Where("c.iccid >= ?", iccidStart) + } + if iccidEnd, ok := filterString(params.Filters, "iccid_end"); ok { + query = query.Where("c.iccid <= ?", iccidEnd) + } + if isReplaced, ok := filterOptionalBool(params.Filters, "is_replaced"); ok { + subQuery := s.db.WithContext(ctx).Table("tb_exchange_order"). + Select("old_asset_id"). + Where("old_asset_type = ? AND status IN ? AND deleted_at IS NULL", + constants.ExchangeAssetTypeIotCard, + []int{constants.ExchangeStatusShipped, constants.ExchangeStatusCompleted}, + ) + if isReplaced { + query = query.Where("c.id IN (?)", subQuery) + } else { + query = query.Where("c.id NOT IN (?)", subQuery) + } + } + if seriesID, ok := filterUint(params.Filters, "series_id"); ok { + query = query.Where("c.series_id = ?", seriesID) + } + if carrierName, ok := filterString(params.Filters, "carrier_name"); ok { + query = query.Where("c.carrier_name LIKE ?", "%"+carrierName+"%") + } + if hasActivePackage, ok := filterOptionalBool(params.Filters, "has_active_package"); ok { + subQuery := s.db.WithContext(ctx).Table("tb_package_usage AS pu"). + Select("1"). + Where("pu.iot_card_id = c.id AND pu.status = ? AND pu.master_usage_id IS NULL AND pu.deleted_at IS NULL", constants.PackageUsageStatusActive) + if hasActivePackage { + query = query.Where("EXISTS (?)", subQuery) + } else { + query = query.Where("NOT EXISTS (?)", subQuery) + } + } + return query +} + +type iotCardExportRow struct { + ICCID string + MSISDN string + DeviceVirtualNo string + CarrierName string + ShopName string + DeviceName string + RealNameStatus int + FirstRealnameAt *time.Time + NetworkStatus int +} + +func formatRealNameVerified(status int) string { + switch status { + case constants.RealNameStatusVerified: + return "是" + case constants.RealNameStatusNotVerified: + return "否" + default: + return "未知" + } +} diff --git a/internal/exporter/query_params.go b/internal/exporter/query_params.go new file mode 100644 index 0000000..e3dacd4 --- /dev/null +++ b/internal/exporter/query_params.go @@ -0,0 +1,56 @@ +package exporter + +import ( + "github.com/bytedance/sonic" + + "github.com/break/junhong_cmp_fiber/internal/model" +) + +// ParseExportParams 从导出任务快照解析数据源执行参数。 +func ParseExportParams(task *model.ExportTask) ExportParams { + params := ExportParams{Filters: make(map[string]any)} + if task == nil { + return params + } + + params.ScopeShopIDs = append(params.ScopeShopIDs, task.ScopeShopIDs...) + params.UserType = task.CreatorUserType + + var raw map[string]any + if len(task.QueryJSON) == 0 { + return params + } + if err := sonic.Unmarshal(task.QueryJSON, &raw); err != nil { + return params + } + + if filters, ok := raw["filters"].(map[string]any); ok { + params.Filters = filters + } + return params +} + +// ResolvedHeadersFromTask 从任务 query_json 读取 dispatch 阶段固定的表头。 +func ResolvedHeadersFromTask(task *model.ExportTask) []string { + if task == nil || len(task.QueryJSON) == 0 { + return nil + } + + var raw map[string]any + if err := sonic.Unmarshal(task.QueryJSON, &raw); err != nil { + return nil + } + + values, ok := raw["resolved_headers"].([]any) + if !ok { + return nil + } + + headers := make([]string, 0, len(values)) + for _, value := range values { + if text, ok := value.(string); ok { + headers = append(headers, text) + } + } + return headers +} diff --git a/internal/exporter/registry.go b/internal/exporter/registry.go index 1ead4c3..491b2d6 100644 --- a/internal/exporter/registry.go +++ b/internal/exporter/registry.go @@ -1,71 +1,54 @@ package exporter import ( - "context" "sort" "gorm.io/gorm" - "github.com/break/junhong_cmp_fiber/internal/model" "github.com/break/junhong_cmp_fiber/pkg/constants" ) -// ShardBoundary 分片边界信息。 -type ShardBoundary struct { - StartID uint64 - EndID uint64 - RowCount int -} - -// SceneStrategy 导出场景策略接口。 -type SceneStrategy interface { - Scene() string - Headers() []string - NextShardBoundary(ctx context.Context, task *model.ExportTask, afterID uint64, limit int) (*ShardBoundary, error) - QueryRows(ctx context.Context, task *model.ExportTask, startID, endID uint64) ([][]string, error) -} - -// Registry 场景策略注册中心。 +// Registry 导出数据源注册中心。 type Registry struct { - strategies map[string]SceneStrategy + sources map[string]DataSource } -// NewRegistry 创建场景策略注册中心。 -func NewRegistry(strategies ...SceneStrategy) *Registry { - m := make(map[string]SceneStrategy, len(strategies)) - for _, strategy := range strategies { - if strategy == nil { +// NewRegistry 创建导出数据源注册中心。 +func NewRegistry(sources ...DataSource) *Registry { + m := make(map[string]DataSource, len(sources)) + for _, source := range sources { + if source == nil { continue } - m[strategy.Scene()] = strategy + m[source.Scene()] = source } - return &Registry{strategies: m} + return &Registry{sources: m} } -// NewDefaultRegistry 创建默认场景策略注册中心。 +// NewDefaultRegistry 创建默认导出数据源注册中心。 func NewDefaultRegistry(db *gorm.DB) *Registry { return NewRegistry( - NewDeviceSceneStrategy(db), - NewIotCardSceneStrategy(db), + NewDeviceDataSource(db), + NewIotCardDataSource(db), ) } -// Get 获取指定场景策略。 -func (r *Registry) Get(scene string) (SceneStrategy, bool) { - strategy, ok := r.strategies[scene] - return strategy, ok +// Get 获取指定场景数据源。 +func (r *Registry) Get(scene string) (DataSource, bool) { + source, ok := r.sources[scene] + return source, ok } // IsSupported 判断场景是否支持。 func (r *Registry) IsSupported(scene string) bool { - _, ok := r.strategies[scene] + _, ok := r.sources[scene] return ok } // Scenes 返回当前支持的场景列表。 func (r *Registry) Scenes() []string { - result := make([]string, 0, len(r.strategies)) - for scene := range r.strategies { + result := make([]string, 0, len(r.sources)) + for scene := range r.sources { result = append(result, scene) } sort.Strings(result) diff --git a/internal/exporter/scope.go b/internal/exporter/scope.go index 69bb1e1..5c64ba7 100644 --- a/internal/exporter/scope.go +++ b/internal/exporter/scope.go @@ -1,26 +1 @@ package exporter - -import ( - "gorm.io/gorm" - - "github.com/break/junhong_cmp_fiber/internal/model" - "github.com/break/junhong_cmp_fiber/pkg/constants" -) - -func applyShopScopeForTask(query *gorm.DB, task *model.ExportTask) *gorm.DB { - if task == nil { - return query.Where("1 = 0") - } - - switch task.CreatorUserType { - case constants.UserTypeSuperAdmin, constants.UserTypePlatform: - return query - case constants.UserTypeAgent: - if len(task.ScopeShopIDs) == 0 { - return query.Where("1 = 0") - } - return query.Where("shop_id IN ?", []uint(task.ScopeShopIDs)) - default: - return query.Where("1 = 0") - } -} diff --git a/internal/model/export_task.go b/internal/model/export_task.go index 8fe0f6f..31418fd 100644 --- a/internal/model/export_task.go +++ b/internal/model/export_task.go @@ -85,11 +85,11 @@ type ExportShardTask struct { gorm.Model BaseModel `gorm:"embedded"` - TaskID uint `gorm:"column:task_id;type:bigint;index:idx_export_shard_task_task_status,priority:1;not null;comment:主任务ID" json:"task_id"` - ShardNo int `gorm:"column:shard_no;type:int;not null;comment:分片序号(从1开始)" json:"shard_no"` - Status int `gorm:"column:status;type:int;index:idx_export_shard_task_task_status,priority:2;not null;default:1;comment:分片状态 1-待处理 2-处理中 3-已成功 4-已失败 5-已取消" json:"status"` - CursorStart uint64 `gorm:"column:cursor_start;type:bigint;not null;default:0;comment:分片游标起点(不含)" json:"cursor_start"` - CursorEnd uint64 `gorm:"column:cursor_end;type:bigint;not null;default:0;comment:分片游标终点(含)" json:"cursor_end"` + TaskID uint `gorm:"column:task_id;type:bigint;index:idx_export_shard_task_task_status,priority:1;not null;comment:主任务ID" json:"task_id"` + ShardNo int `gorm:"column:shard_no;type:int;not null;comment:分片序号(从1开始)" json:"shard_no"` + Status int `gorm:"column:status;type:int;index:idx_export_shard_task_task_status,priority:2;not null;default:1;comment:分片状态 1-待处理 2-处理中 3-已成功 4-已失败 5-已取消" json:"status"` + ShardOffset int `gorm:"column:shard_offset;type:bigint;not null;default:0;comment:分片偏移量(offset)" json:"shard_offset"` + ShardLimit int `gorm:"column:shard_limit;type:int;not null;default:0;comment:分片行数上限(limit)" json:"shard_limit"` RowCount int `gorm:"column:row_count;type:int;not null;default:0;comment:分片行数" json:"row_count"` ProcessedRows int `gorm:"column:processed_rows;type:int;not null;default:0;comment:分片已处理行数" json:"processed_rows"` diff --git a/internal/store/postgres/export_task_store.go b/internal/store/postgres/export_task_store.go index 9775016..17cf15c 100644 --- a/internal/store/postgres/export_task_store.go +++ b/internal/store/postgres/export_task_store.go @@ -5,7 +5,9 @@ import ( "fmt" "time" + "github.com/bytedance/sonic" "github.com/redis/go-redis/v9" + "gorm.io/datatypes" "gorm.io/gorm" "gorm.io/gorm/clause" @@ -124,7 +126,7 @@ func (s *ExportTaskStore) TryMarkProcessingIfPending(ctx context.Context, taskID } // UpdateDispatchPlan 更新 dispatch 阶段计算出的总量信息。 -func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, totalRows, totalShards uint, updater uint) error { +func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, totalRows, totalShards int, updater uint) error { now := time.Now() return s.db.WithContext(ctx). Model(&model.ExportTask{}). @@ -138,6 +140,40 @@ func (s *ExportTaskStore) UpdateDispatchPlan(ctx context.Context, taskID uint, t }).Error } +// UpdateResolvedHeaders 固定导出任务运行期表头。 +func (s *ExportTaskStore) UpdateResolvedHeaders(ctx context.Context, taskID uint, headers []string, updater uint) error { + var task model.ExportTask + if err := s.db.WithContext(ctx). + Select("id", "query_json"). + Where("id = ?", taskID). + First(&task).Error; err != nil { + return err + } + + raw := make(map[string]any) + if len(task.QueryJSON) > 0 { + if err := sonic.Unmarshal(task.QueryJSON, &raw); err != nil { + return err + } + } + raw["resolved_headers"] = headers + + bytes, err := sonic.Marshal(raw) + if err != nil { + return err + } + + now := time.Now() + return s.db.WithContext(ctx). + Model(&model.ExportTask{}). + Where("id = ?", taskID). + Updates(map[string]any{ + "query_json": datatypes.JSON(bytes), + "updated_at": now, + "updater": updater, + }).Error +} + // IsCancelRequested 查询是否已请求取消。 func (s *ExportTaskStore) IsCancelRequested(ctx context.Context, taskID uint) (bool, error) { var value bool diff --git a/internal/task/export_common.go b/internal/task/export_common.go index 3741c78..c4bca61 100644 --- a/internal/task/export_common.go +++ b/internal/task/export_common.go @@ -5,6 +5,7 @@ import ( "encoding/csv" "errors" "fmt" + "io" "os" "path/filepath" "time" @@ -82,18 +83,18 @@ func isFinalRetry(ctx context.Context) bool { return retryCount >= maxRetry } -func writeExportFile(format string, headers []string, rows [][]string) (string, int64, error) { +func writeExportFile(format string, headers []string, rows [][]string, withoutHeader bool) (string, int64, error) { switch format { case constants.ExportTaskFormatCSV: - return writeCSVFile(headers, rows) + return writeCSVFile(headers, rows, withoutHeader) case constants.ExportTaskFormatXLSX: - return writeXLSXFile(headers, rows) + return writeXLSXFile(headers, rows, withoutHeader) default: return "", 0, fmt.Errorf("不支持的导出格式: %s", format) } } -func writeCSVFile(headers []string, rows [][]string) (string, int64, error) { +func writeCSVFile(headers []string, rows [][]string, withoutHeader bool) (string, int64, error) { file, err := os.CreateTemp("", "export-*.csv") if err != nil { return "", 0, err @@ -101,8 +102,10 @@ func writeCSVFile(headers []string, rows [][]string) (string, int64, error) { defer file.Close() writer := csv.NewWriter(file) - if err := writer.Write(headers); err != nil { - return "", 0, err + if !withoutHeader { + if err := writer.Write(headers); err != nil { + return "", 0, err + } } for _, row := range rows { if err := writer.Write(row); err != nil { @@ -125,7 +128,7 @@ func writeCSVFile(headers []string, rows [][]string) (string, int64, error) { return file.Name(), stat.Size(), nil } -func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) { +func writeXLSXFile(headers []string, rows [][]string, withoutHeader bool) (string, int64, error) { file := excelize.NewFile() sheet := file.GetSheetName(file.GetActiveSheetIndex()) @@ -134,16 +137,20 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) { 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 + rowIndex := 1 + if !withoutHeader { + 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 + } + rowIndex++ } - for i, row := range rows { - axis, err := excelize.CoordinatesToCellName(1, i+2) + for _, row := range rows { + axis, err := excelize.CoordinatesToCellName(1, rowIndex) if err != nil { return "", 0, err } @@ -154,6 +161,7 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) { if err := streamWriter.SetRow(axis, values); err != nil { return "", 0, err } + rowIndex++ } if err := streamWriter.Flush(); err != nil { @@ -175,6 +183,90 @@ func writeXLSXFile(headers []string, rows [][]string) (string, int64, error) { return path, stat.Size(), nil } +func writeXLSXFromCSV(csvPath string) (string, int64, error) { + source, err := os.Open(csvPath) + if err != nil { + return "", 0, err + } + defer source.Close() + + file := excelize.NewFile() + sheet := file.GetSheetName(file.GetActiveSheetIndex()) + streamWriter, err := file.NewStreamWriter(sheet) + if err != nil { + return "", 0, err + } + + reader := csv.NewReader(source) + rowIndex := 1 + for { + row, err := reader.Read() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return "", 0, err + } + + values := make([]interface{}, 0, len(row)) + for _, item := range row { + values = append(values, item) + } + axis, err := excelize.CoordinatesToCellName(1, rowIndex) + if err != nil { + return "", 0, err + } + if err := streamWriter.SetRow(axis, values); err != nil { + return "", 0, err + } + rowIndex++ + } + + 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 alignRowsToHeaders(rows [][]string, headers []string) [][]string { + if len(headers) == 0 { + return rows + } + + aligned := make([][]string, 0, len(rows)) + for _, row := range rows { + current := make([]string, len(headers)) + copy(current, row) + aligned = append(aligned, current) + } + return aligned +} + +func hasMisalignedRows(rows [][]string, headers []string) bool { + if len(headers) == 0 { + return false + } + for _, row := range rows { + if len(row) != len(headers) { + return true + } + } + return false +} + func uploadExportFile(ctx context.Context, storageSvc *storage.Service, localPath, fileName string) (string, error) { if storageSvc == nil || storageSvc.Provider() == nil { return "", errors.New("对象存储服务未配置") diff --git a/internal/task/export_dispatch.go b/internal/task/export_dispatch.go index 372928f..9b1648d 100644 --- a/internal/task/export_dispatch.go +++ b/internal/task/export_dispatch.go @@ -119,14 +119,22 @@ func (h *ExportDispatchHandler) HandleExportDispatch(ctx context.Context, task * return asynq.SkipRetry } - shards, totalRows, err := h.buildShards(ctx, exportTask, strategy, updater) + params := exporter.ParseExportParams(exportTask) + headers, err := strategy.Headers(ctx, params) + if err != nil { + _ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "解析导出表头失败") + h.logger.Error("解析导出表头失败", zap.Uint("task_id", exportTask.ID), zap.Error(err)) + return err + } + + shards, totalRows, err := h.buildShards(ctx, strategy, params, exportTask.CancelRequested, exportTask.ID, updater) if err != nil { _ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "构建导出分片失败") h.logger.Error("构建导出分片失败", zap.Uint("task_id", exportTask.ID), zap.Error(err)) return err } - if err := h.persistShardsAndPlan(ctx, exportTask.ID, updater, shards, totalRows); err != nil { + if err := h.persistShardsAndPlan(ctx, exportTask.ID, updater, shards, totalRows, headers); err != nil { _ = h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "保存导出分片失败") return err } @@ -141,58 +149,53 @@ func (h *ExportDispatchHandler) HandleExportDispatch(ctx context.Context, task * h.logger.Info("导出 dispatch 完成", zap.Uint("task_id", exportTask.ID), - zap.Uint("total_rows", totalRows), + zap.Int("total_rows", totalRows), zap.Int("total_shards", len(shards)), ) return nil } -func (h *ExportDispatchHandler) buildShards(ctx context.Context, task *model.ExportTask, strategy exporter.SceneStrategy, updater uint) ([]*model.ExportShardTask, uint, error) { +func (h *ExportDispatchHandler) buildShards(ctx context.Context, source exporter.DataSource, params exporter.ExportParams, cancelRequested bool, taskID uint, updater uint) ([]*model.ExportShardTask, int, error) { shards := make([]*model.ExportShardTask, 0) - totalRows := uint(0) - afterID := uint64(0) + if cancelRequested { + return shards, 0, nil + } + + totalRows, err := source.Count(ctx, params) + if err != nil { + return nil, 0, err + } + if totalRows <= 0 { + return shards, 0, nil + } + shardNo := 1 - - for { - if task.CancelRequested { - break + for offset := 0; offset < totalRows; offset += constants.ExportDefaultShardSize { + limit := constants.ExportDefaultShardSize + if remaining := totalRows - offset; remaining < limit { + limit = remaining } - boundary, err := strategy.NextShardBoundary(ctx, task, afterID, constants.ExportDefaultShardSize) - if err != nil { - return nil, 0, err - } - if boundary == nil { - break - } - - rowCount := boundary.RowCount - if rowCount < 0 { - rowCount = 0 - } - - totalRows += uint(rowCount) shard := &model.ExportShardTask{ - TaskID: task.ID, + TaskID: taskID, ShardNo: shardNo, Status: constants.ExportShardStatusPending, - CursorStart: boundary.StartID, - CursorEnd: boundary.EndID, - RowCount: rowCount, + ShardOffset: offset, + ShardLimit: limit, + RowCount: limit, } shard.Creator = updater shard.Updater = updater shards = append(shards, shard) - afterID = boundary.EndID shardNo++ } return shards, totalRows, nil } -func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows uint) error { +func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID uint, updater uint, shards []*model.ExportShardTask, totalRows int, headers []string) error { return h.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { txShardStore := postgres.NewExportShardTaskStore(tx, h.redis) txTaskStore := postgres.NewExportTaskStore(tx, h.redis) @@ -200,7 +203,10 @@ func (h *ExportDispatchHandler) persistShardsAndPlan(ctx context.Context, taskID if err := txShardStore.CreateBatch(ctx, shards); err != nil { return err } - if err := txTaskStore.UpdateDispatchPlan(ctx, taskID, totalRows, uint(len(shards)), updater); err != nil { + if err := txTaskStore.UpdateResolvedHeaders(ctx, taskID, headers, updater); err != nil { + return err + } + if err := txTaskStore.UpdateDispatchPlan(ctx, taskID, totalRows, len(shards), updater); err != nil { return err } return nil diff --git a/internal/task/export_finalize.go b/internal/task/export_finalize.go index 09c5eda..5dfc0af 100644 --- a/internal/task/export_finalize.go +++ b/internal/task/export_finalize.go @@ -2,7 +2,10 @@ package task import ( "context" + "encoding/csv" + "errors" "fmt" + "io" "os" "github.com/hibiken/asynq" @@ -10,6 +13,7 @@ import ( "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" @@ -102,8 +106,7 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task * return nil } - strategy, ok := h.sceneRegistry.Get(exportTask.Scene) - if !ok { + if !h.sceneRegistry.IsSupported(exportTask.Scene) { if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "导出场景未注册"); err != nil { return err } @@ -115,23 +118,22 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task * 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...) + headers := exporter.ResolvedHeadersFromTask(exportTask) + if len(headers) == 0 { + return h.handleFinalizeError(ctx, exportTask.ID, updater, "读取导出表头失败", errors.New("resolved_headers 为空")) } - localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows) + mergedCSVPath, cleanupCSV, err := h.mergeShardCSVFiles(ctx, headers, shards) if err != nil { return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err) } - defer func() { _ = os.Remove(localPath) }() + 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) @@ -146,6 +148,104 @@ func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task * 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()) diff --git a/internal/task/export_shard.go b/internal/task/export_shard.go index 88d308b..3115616 100644 --- a/internal/task/export_shard.go +++ b/internal/task/export_shard.go @@ -107,24 +107,38 @@ func (h *ExportShardHandler) HandleExportShard(ctx context.Context, task *asynq. } } - strategy, ok := h.sceneRegistry.Get(exportTask.Scene) + source, ok := h.sceneRegistry.Get(exportTask.Scene) if !ok { _ = h.markShardFinalFailed(ctx, exportTask.ID, shardTask.ID, updater, "导出场景未注册") return asynq.SkipRetry } - rows, err := strategy.QueryRows(ctx, exportTask, shardTask.CursorStart, shardTask.CursorEnd) + headers := exporter.ResolvedHeadersFromTask(exportTask) + if len(headers) == 0 { + return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "读取导出表头失败", fmt.Errorf("resolved_headers 为空")) + } + + params := exporter.ParseExportParams(exportTask) + rows, err := source.Fetch(ctx, params, shardTask.ShardOffset, shardTask.ShardLimit) if err != nil { return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "查询分片数据失败", err) } + if hasMisalignedRows(rows, headers) { + h.logger.Warn("导出分片数据列数与表头不一致,已按固定表头对齐", + zap.Uint("task_id", exportTask.ID), + zap.Uint("shard_id", shardTask.ID), + zap.Int("header_count", len(headers)), + ) + } + rows = alignRowsToHeaders(rows, headers) - localPath, fileSize, err := writeExportFile(exportTask.Format, strategy.Headers(), rows) + localPath, fileSize, err := writeExportFile(constants.ExportTaskFormatCSV, headers, rows, true) if err != nil { return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "生成分片文件失败", err) } defer func() { _ = os.Remove(localPath) }() - fileName := fmt.Sprintf("%s-shard-%04d.%s", exportTask.TaskNo, shardTask.ShardNo, exportTask.Format) + fileName := fmt.Sprintf("%s-shard-%04d.csv.shard", exportTask.TaskNo, shardTask.ShardNo) fileKey, err := uploadExportFile(ctx, h.storageSvc, localPath, fileName) if err != nil { return h.handleShardError(ctx, exportTask.ID, shardTask.ID, updater, "上传分片文件失败", err) diff --git a/migrations/000153_alter_export_shard_task_offset.down.sql b/migrations/000153_alter_export_shard_task_offset.down.sql new file mode 100644 index 0000000..ac7924b --- /dev/null +++ b/migrations/000153_alter_export_shard_task_offset.down.sql @@ -0,0 +1,4 @@ +-- 回滚导出分片任务 offset/limit 分片字段 +ALTER TABLE tb_export_shard_task + DROP COLUMN IF EXISTS shard_limit, + DROP COLUMN IF EXISTS shard_offset; diff --git a/migrations/000153_alter_export_shard_task_offset.up.sql b/migrations/000153_alter_export_shard_task_offset.up.sql new file mode 100644 index 0000000..9cb150e --- /dev/null +++ b/migrations/000153_alter_export_shard_task_offset.up.sql @@ -0,0 +1,7 @@ +-- 为导出分片任务补充 offset/limit 分片字段 +ALTER TABLE tb_export_shard_task + ADD COLUMN IF NOT EXISTS shard_offset BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS shard_limit INT NOT NULL DEFAULT 0; + +COMMENT ON COLUMN tb_export_shard_task.shard_offset IS '分片偏移量(offset)'; +COMMENT ON COLUMN tb_export_shard_task.shard_limit IS '分片行数上限(limit)'; diff --git a/pkg/storage/s3.go b/pkg/storage/s3.go index 2f26f93..6763eea 100644 --- a/pkg/storage/s3.go +++ b/pkg/storage/s3.go @@ -73,7 +73,7 @@ func (p *S3Provider) Upload(ctx context.Context, key string, reader io.Reader, c return nil } -func (p *S3Provider) Download(ctx context.Context, key string, writer io.Writer) error { +func (p *S3Provider) Download(ctx context.Context, key string) (io.ReadCloser, error) { input := &s3.GetObjectInput{ Bucket: aws.String(p.bucket), Key: aws.String(key), @@ -82,17 +82,11 @@ func (p *S3Provider) Download(ctx context.Context, key string, writer io.Writer) result, err := p.client.GetObjectWithContext(ctx, input) if err != nil { if strings.Contains(err.Error(), "NoSuchKey") { - return fmt.Errorf("文件不存在: %s", key) + return nil, fmt.Errorf("文件不存在: %s", key) } - return fmt.Errorf("下载文件失败: %w", err) + return nil, fmt.Errorf("下载文件失败: %w", err) } - defer result.Body.Close() - - _, err = io.Copy(writer, result.Body) - if err != nil { - return fmt.Errorf("写入文件内容失败: %w", err) - } - return nil + return result.Body, nil } func (p *S3Provider) DownloadToTemp(ctx context.Context, key string) (string, func(), error) { @@ -111,11 +105,19 @@ func (p *S3Provider) DownloadToTemp(ctx context.Context, key string) (string, fu _ = os.Remove(tempPath) } - if err := p.Download(ctx, key, tempFile); err != nil { + reader, err := p.Download(ctx, key) + if err != nil { tempFile.Close() cleanup() return "", nil, err } + defer reader.Close() + + if _, err := io.Copy(tempFile, reader); err != nil { + tempFile.Close() + cleanup() + return "", nil, fmt.Errorf("写入文件内容失败: %w", err) + } if err := tempFile.Close(); err != nil { cleanup() diff --git a/pkg/storage/storage.go b/pkg/storage/storage.go index e00a2cd..1128173 100644 --- a/pkg/storage/storage.go +++ b/pkg/storage/storage.go @@ -8,7 +8,7 @@ import ( type Provider interface { Upload(ctx context.Context, key string, reader io.Reader, contentType string) error - Download(ctx context.Context, key string, writer io.Writer) error + Download(ctx context.Context, key string) (io.ReadCloser, error) DownloadToTemp(ctx context.Context, key string) (localPath string, cleanup func(), err error) Delete(ctx context.Context, key string) error Exists(ctx context.Context, key string) (bool, error)