junhong_cmp_fiber/internal/exporter/registry.go

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 场景策略注册中心。
type Registry struct {
	strategies map[string]SceneStrategy
}

// NewRegistry 创建场景策略注册中心。
func NewRegistry(strategies ...SceneStrategy) *Registry {
	m := make(map[string]SceneStrategy, len(strategies))
	for _, strategy := range strategies {
		if strategy == nil {
			continue
		}
		m[strategy.Scene()] = strategy
	}
	return &Registry{strategies: m}
}

// NewDefaultRegistry 创建默认场景策略注册中心。
func NewDefaultRegistry(db *gorm.DB) *Registry {
	return NewRegistry(
		NewDeviceSceneStrategy(db),
		NewIotCardSceneStrategy(db),
	)
}

// Get 获取指定场景策略。
func (r *Registry) Get(scene string) (SceneStrategy, bool) {
	strategy, ok := r.strategies[scene]
	return strategy, ok
}

// IsSupported 判断场景是否支持。
func (r *Registry) IsSupported(scene string) bool {
	_, ok := r.strategies[scene]
	return ok
}

// Scenes 返回当前支持的场景列表。
func (r *Registry) Scenes() []string {
	result := make([]string, 0, len(r.strategies))
	for scene := range r.strategies {
		result = append(result, scene)
	}
	sort.Strings(result)
	return result
}

// IsSupportedScene 判断是否为受支持的场景。
func IsSupportedScene(scene string) bool {
	switch scene {
	case constants.ExportTaskSceneDevice, constants.ExportTaskSceneIotCard:
		return true
	default:
		return false
	}
}