重构seeder
This commit is contained in:
@@ -7,49 +7,81 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"git.huangwc.com/pig/pig-farm-controller/internal/infra/database/seeder"
|
||||
"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// SeederFunc 定义了处理一种特定类型预设数据文件的函数签名。
|
||||
type SeederFunc func(ctx context.Context, tx *gorm.DB, jsonData []byte) error
|
||||
|
||||
// isTableEmpty 检查给定模型对应的数据库表是否为空。
|
||||
func isTableEmpty(tx *gorm.DB, model interface{}) (bool, error) {
|
||||
var count int64
|
||||
if err := tx.Model(model).Count(&count).Error; err != nil {
|
||||
return false, fmt.Errorf("查询表记录数失败: %w", err)
|
||||
}
|
||||
return count == 0, nil
|
||||
// Seeder 定义了数据播种器的通用接口。
|
||||
// 每个播种器负责处理一种特定类型的预设数据。
|
||||
type Seeder interface {
|
||||
// Type 返回此播种器能处理的唯一类型标识符，
|
||||
// 这个标识符应与预设 JSON 文件中的 "type" 字段相匹配。
|
||||
Type() string
|
||||
// IsRequired 标记此播种器对应的预设文件是否为必需的。
|
||||
// 如果为 true 且对应的 .json 文件不存在，整个播种过程将失败并报错。
|
||||
IsRequired() bool
|
||||
// Seed 执行具体的数据播种逻辑。
|
||||
// 它接收数据库事务、上下文和从 JSON 文件读取的原始字节数据。
|
||||
Seed(ctx context.Context, tx *gorm.DB, jsonData []byte) error
|
||||
}
|
||||
|
||||
// SeedFromPreset 是一个通用的数据播种函数。
|
||||
// 它会读取指定目录下的所有 .json 文件，并根据文件内容中的 "type" 字段进行分发。
|
||||
// 同时，它会校验所有必需的预设类型是否都已成功加载。
|
||||
func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string) error {
|
||||
// 它会读取指定目录下的所有 .json 文件，并根据文件内容中的 "type" 字段，
|
||||
// 将其分发给由外部注入的相应 Seeder 进行处理。
|
||||
// Seeder 的处理顺序由其在注入切片中的顺序决定。
|
||||
func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string, seeders []Seeder) error {
|
||||
seedCtx, logger := logs.Trace(ctx, ctx, "SeedFromPreset")
|
||||
|
||||
// 定义必须存在的预设数据类型及其处理顺序
|
||||
// 确保 "nutrient" 在 "pig_nutrient_requirements" 之前处理，因为后者依赖于前者。
|
||||
processingOrder := []string{"nutrient", "pig_nutrient_requirements"}
|
||||
requiredTypes := make(map[string]bool)
|
||||
for _, t := range processingOrder {
|
||||
requiredTypes[t] = true
|
||||
// --- 步骤 1: 校验注入的 Seeders 是否有重复类型 ---
|
||||
tempSeederTypes := make(map[string]bool)
|
||||
for _, s := range seeders {
|
||||
if tempSeederTypes[s.Type()] {
|
||||
return fmt.Errorf("播种器初始化失败: 存在重复的 Seeder 类型 '%s'", s.Type())
|
||||
}
|
||||
tempSeederTypes[s.Type()] = true
|
||||
}
|
||||
|
||||
processedTypes := make(map[string]bool)
|
||||
typeToFileMap := make(map[string]string) // 用于检测重复的 type，并存储每个 type 对应的文件路径
|
||||
groupedFiles := make(map[string][][]byte) // 按 type 分组存储 jsonData
|
||||
// --- 步骤 2: 读取并分组所有预设文件 ---
|
||||
groupedFiles := make(map[string][][]byte)
|
||||
typeToFileMap := make(map[string]string)
|
||||
|
||||
files, err := os.ReadDir(presetDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) { // 目录不存在
|
||||
// 检查是否有必需的 Seeder，如果目录不存在但需要文件，则报错
|
||||
var requiredSeederTypes []string
|
||||
for _, s := range seeders {
|
||||
if s.IsRequired() {
|
||||
requiredSeederTypes = append(requiredSeederTypes, s.Type())
|
||||
}
|
||||
}
|
||||
if len(requiredSeederTypes) > 0 {
|
||||
return fmt.Errorf("预设数据校验失败: 预设目录 '%s' 不存在, 但系统需要以下必需类型: [%s]", presetDir, strings.Join(requiredSeederTypes, ", "))
|
||||
}
|
||||
logger.Warnf("预设数据目录 '%s' 不存在，跳过播种。", presetDir)
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("读取预设数据目录 '%s' 失败: %w", presetDir, err)
|
||||
}
|
||||
|
||||
// 第一阶段：读取所有文件并按 type 分组
|
||||
// 处理目录为空的情况
|
||||
if len(files) == 0 {
|
||||
var requiredSeederTypes []string
|
||||
for _, s := range seeders {
|
||||
if s.IsRequired() {
|
||||
requiredSeederTypes = append(requiredSeederTypes, s.Type())
|
||||
}
|
||||
}
|
||||
if len(requiredSeederTypes) > 0 {
|
||||
return fmt.Errorf("预设数据校验失败: 预设目录 '%s' 为空, 但系统需要以下必需类型: [%s]", presetDir, strings.Join(requiredSeederTypes, ", "))
|
||||
}
|
||||
logger.Warnf("预设数据目录 '%s' 为空，跳过播种。", presetDir)
|
||||
return nil
|
||||
}
|
||||
|
||||
// 读取所有文件并按 type 分组
|
||||
for _, file := range files {
|
||||
if filepath.Ext(file.Name()) != ".json" {
|
||||
continue
|
||||
@@ -63,7 +95,7 @@ func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string) error {
|
||||
|
||||
dataType := gjson.GetBytes(jsonData, "type")
|
||||
if !dataType.Exists() {
|
||||
logger.Warnf("警告: 文件 '%s' 中缺少 'type' 字段，已跳过", filePath)
|
||||
logger.Warnf("警告: 文件 '%s' 中缺少 'type' 字段，已跳过。", filePath)
|
||||
continue
|
||||
}
|
||||
dataTypeStr := dataType.String()
|
||||
@@ -77,50 +109,37 @@ func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string) error {
|
||||
groupedFiles[dataTypeStr] = append(groupedFiles[dataTypeStr], jsonData)
|
||||
}
|
||||
|
||||
// 第二阶段：按照预定义顺序处理分组后的数据
|
||||
// --- 步骤 3: 在事务中按注入顺序执行播种，并采用快速失败策略 ---
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
for _, dataTypeStr := range processingOrder {
|
||||
jsonDatas, ok := groupedFiles[dataTypeStr]
|
||||
if !ok {
|
||||
// 如果是必需类型但没有找到文件，则报错
|
||||
if requiredTypes[dataTypeStr] {
|
||||
logger.Info("开始执行数据播种事务...")
|
||||
|
||||
// 直接遍历注入的 seeders，顺序由调用方保证
|
||||
for _, seeder := range seeders {
|
||||
dataTypeStr := seeder.Type()
|
||||
jsonDatas, fileExists := groupedFiles[dataTypeStr]
|
||||
|
||||
if fileExists {
|
||||
// --- 文件存在，执行播种 ---
|
||||
logger.Infof("正在使用播种器处理类型: '%s'...", dataTypeStr)
|
||||
for _, jsonData := range jsonDatas {
|
||||
originalFilePath := typeToFileMap[dataTypeStr]
|
||||
if err := seeder.Seed(seedCtx, tx, jsonData); err != nil {
|
||||
return fmt.Errorf("处理文件 (type: %s, path: %s) 时发生错误: %w", dataTypeStr, originalFilePath, err)
|
||||
}
|
||||
}
|
||||
logger.Infof("类型 '%s' 处理完成。", dataTypeStr)
|
||||
} else {
|
||||
// --- 文件不存在，检查是否必需 ---
|
||||
if seeder.IsRequired() {
|
||||
// 快速失败：如果是必需的但文件不存在，立即报错并回滚事务
|
||||
return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: '%s'", dataTypeStr)
|
||||
}
|
||||
continue // 非必需类型，跳过
|
||||
}
|
||||
|
||||
var seederFunc SeederFunc
|
||||
switch dataTypeStr {
|
||||
case "nutrient":
|
||||
seederFunc = seeder.SeedNutrients
|
||||
case "pig_nutrient_requirements":
|
||||
seederFunc = seeder.SeedPigNutrientRequirements
|
||||
default:
|
||||
logger.Warnf("警告: 存在未知的 type: '%s'，已跳过", dataTypeStr)
|
||||
continue
|
||||
}
|
||||
|
||||
for _, jsonData := range jsonDatas {
|
||||
// 获取原始文件路径用于错误报告
|
||||
originalFilePath := typeToFileMap[dataTypeStr]
|
||||
if err := seederFunc(seedCtx, tx, jsonData); err != nil {
|
||||
return fmt.Errorf("处理文件 (type: %s, path: %s) 时发生错误: %w", dataTypeStr, originalFilePath, err)
|
||||
}
|
||||
}
|
||||
processedTypes[dataTypeStr] = true
|
||||
}
|
||||
|
||||
// 校验所有必需的类型是否都已处理
|
||||
var missingTypes []string
|
||||
for reqType := range requiredTypes {
|
||||
if !processedTypes[reqType] {
|
||||
missingTypes = append(missingTypes, reqType)
|
||||
// 如果不是必需的，则只记录日志
|
||||
logger.Infof("未找到可选类型为 '%s' 的预设文件，跳过该播种器。", dataTypeStr)
|
||||
}
|
||||
}
|
||||
if len(missingTypes) > 0 {
|
||||
return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: [%s]", strings.Join(missingTypes, ", "))
|
||||
}
|
||||
|
||||
return nil // 提交事务
|
||||
logger.Info("数据播种事务成功完成。")
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user