导出系统
All checks were successful
构建并部署到测试环境(无 SSH) / build-and-deploy (push) Successful in 7m50s

This commit is contained in:
Break
2026-06-15 15:28:29 +08:00
parent 2f0b24ce88
commit 7ec84fbc0f
17 changed files with 939 additions and 260 deletions

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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)
}

View File

@@ -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 "未知"
}
}

View File

@@ -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
}

View File

@@ -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)

View File

@@ -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")
}
}

View File

@@ -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"`

View File

@@ -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

View File

@@ -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("对象存储服务未配置")

View File

@@ -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

View File

@@ -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())

View File

@@ -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)

View File

@@ -0,0 +1,4 @@
-- 回滚导出分片任务 offset/limit 分片字段
ALTER TABLE tb_export_shard_task
DROP COLUMN IF EXISTS shard_limit,
DROP COLUMN IF EXISTS shard_offset;

View File

@@ -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)';

View File

@@ -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()

View File

@@ -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)