Files
pig-farm-controller/internal/infra/database/seeder.go
2025-11-21 18:36:02 +08:00

521 lines
18 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package database
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
"github.com/tidwall/gjson"
"gorm.io/gorm"
)
// SeederFunc 定义了处理一种特定类型预设数据文件的函数签名。
type SeederFunc func(tx *gorm.DB, jsonData []byte) error
// SeedFromPreset 是一个通用的数据播种函数。
// 它会读取指定目录下的所有 .json 文件,并根据文件内容中的 "type" 字段进行分发。
// 同时,它会校验所有必需的预设类型是否都已成功加载。
func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string) error {
logger := logs.TraceLogger(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
}
processedTypes := make(map[string]bool)
typeToFileMap := make(map[string]string) // 用于检测重复的 type并存储每个 type 对应的文件路径
groupedFiles := make(map[string][][]byte) // 按 type 分组存储 jsonData
files, err := os.ReadDir(presetDir)
if err != nil {
return fmt.Errorf("读取预设数据目录 '%s' 失败: %w", presetDir, err)
}
// 第一阶段:读取所有文件并按 type 分组
for _, file := range files {
if filepath.Ext(file.Name()) != ".json" {
continue
}
filePath := filepath.Join(presetDir, file.Name())
jsonData, err := os.ReadFile(filePath)
if err != nil {
return fmt.Errorf("读取文件 '%s' 失败: %w", filePath, err)
}
dataType := gjson.GetBytes(jsonData, "type")
if !dataType.Exists() {
logger.Warnf("警告: 文件 '%s' 中缺少 'type' 字段,已跳过", filePath)
continue
}
dataTypeStr := dataType.String()
// 检查是否存在重复的 type
if existingFile, found := typeToFileMap[dataTypeStr]; found {
return fmt.Errorf("预设数据校验失败: type '%s' 在文件 '%s' 和 '%s' 中重复定义", dataTypeStr, existingFile, filePath)
}
typeToFileMap[dataTypeStr] = filePath // 记录该 type 对应的文件路径
groupedFiles[dataTypeStr] = append(groupedFiles[dataTypeStr], jsonData)
}
// 第二阶段:按照预定义顺序处理分组后的数据
return db.Transaction(func(tx *gorm.DB) error {
for _, dataTypeStr := range processingOrder {
jsonDatas, ok := groupedFiles[dataTypeStr]
if !ok {
// 如果是必需类型但没有找到文件,则报错
if requiredTypes[dataTypeStr] {
return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: '%s'", dataTypeStr)
}
continue // 非必需类型,跳过
}
var seederFunc SeederFunc
switch dataTypeStr {
case "nutrient":
seederFunc = seedNutrients
case "pig_nutrient_requirements":
seederFunc = seedPigNutrientRequirements
default:
logger.Warnf("警告: 存在未知的 type: '%s',已跳过", dataTypeStr)
continue
}
for _, jsonData := range jsonDatas {
// 获取原始文件路径用于错误报告
originalFilePath := typeToFileMap[dataTypeStr]
if err := seederFunc(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)
}
}
if len(missingTypes) > 0 {
return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: [%s]", strings.Join(missingTypes, ", "))
}
return nil // 提交事务
})
}
// seedNutrients 先严格校验JSON源文件然后以“有则跳过”的模式播种数据。
func seedNutrients(tx *gorm.DB, jsonData []byte) error {
// 1. 严格校验JSON文件检查内部重复键
if err := validateAndParseNutrientJSON(jsonData); err != nil {
return fmt.Errorf("JSON源文件校验失败: %w", err)
}
// 2. 解析简介信息
descriptionsNode := gjson.GetBytes(jsonData, "descriptions")
rawMaterialDescriptions := make(map[string]string)
nutrientDescriptions := make(map[string]string)
if descriptionsNode.Exists() {
descriptionsNode.Get("raw_materials").ForEach(func(key, value gjson.Result) bool {
rawMaterialDescriptions[key.String()] = value.String()
return true
})
descriptionsNode.Get("nutrients").ForEach(func(key, value gjson.Result) bool {
nutrientDescriptions[key.String()] = value.String()
return true
})
}
// 3. 将通过校验的、干净的数据写入数据库
dataNode := gjson.GetBytes(jsonData, "data")
var err error // 用于捕获 ForEach 内部的错误
dataNode.ForEach(func(rawMaterialKey, rawMaterialValue gjson.Result) bool {
rawMaterialName := rawMaterialKey.String()
var rawMaterial models.RawMaterial
// 将 Description 放入 Create 对象中
err = tx.Where(models.RawMaterial{Name: rawMaterialName}).
FirstOrCreate(&rawMaterial, models.RawMaterial{
Name: rawMaterialName,
Description: rawMaterialDescriptions[rawMaterialName],
}).Error
if err != nil {
// 返回 false 停止 ForEach 遍历
return false
}
rawMaterialValue.ForEach(func(nutrientKey, nutrientValue gjson.Result) bool {
nutrientName := nutrientKey.String()
value := float32(nutrientValue.Float())
var nutrient models.Nutrient
// 将 Description 放入 Create 对象中
err = tx.Where(models.Nutrient{Name: nutrientName}).
FirstOrCreate(&nutrient, models.Nutrient{
Name: nutrientName,
Description: nutrientDescriptions[nutrientName],
}).Error
if err != nil {
// 返回 false 停止 ForEach 遍历
return false
}
linkData := models.RawMaterialNutrient{
RawMaterialID: rawMaterial.ID,
NutrientID: nutrient.ID,
}
// 使用 FirstOrCreate 确保关联的唯一性
err = tx.Where(linkData).FirstOrCreate(&linkData, models.RawMaterialNutrient{
RawMaterialID: linkData.RawMaterialID,
NutrientID: linkData.NutrientID,
Value: value,
}).Error
if err != nil {
// 返回 false 停止 ForEach 遍历
return false
}
return true
})
return err == nil // 如果内部遍历有错误,则停止外部遍历
})
return err // 返回捕获到的错误
}
// seedPigNutrientRequirements 先严格校验JSON源文件然后以“有则跳过”的模式播种数据。
func seedPigNutrientRequirements(tx *gorm.DB, jsonData []byte) error {
// 1. 严格校验JSON文件检查内部重复键
if err := validateAndParsePigNutrientRequirementJSON(jsonData); err != nil {
return fmt.Errorf("JSON源文件校验失败: %w", err)
}
// 2. 解析简介信息
descriptionsNode := gjson.GetBytes(jsonData, "descriptions")
pigBreedDescriptions := make(map[string]models.PigBreed)
pigAgeStageDescriptions := make(map[string]models.PigAgeStage)
pigTypeDescriptions := make(map[string]map[string]models.PigType)
if descriptionsNode.Exists() {
// 解析 pig_breeds 描述
descriptionsNode.Get("pig_breeds").ForEach(func(key, value gjson.Result) bool {
var pb models.PigBreed
pb.Name = key.String()
pb.Description = value.Get("description").String()
pb.ParentInfo = value.Get("parent_info").String()
pb.AppearanceFeatures = value.Get("appearance_features").String()
pb.BreedAdvantages = value.Get("breed_advantages").String()
pb.BreedDisadvantages = value.Get("breed_disadvantages").String()
pigBreedDescriptions[key.String()] = pb
return true
})
// 解析 pig_age_stages 描述
descriptionsNode.Get("pig_age_stages").ForEach(func(key, value gjson.Result) bool {
var pas models.PigAgeStage
pas.Name = key.String()
pas.Description = value.String()
pigAgeStageDescriptions[key.String()] = pas
return true
})
// 解析 pig_breed_age_stages (PigType) 描述
descriptionsNode.Get("pig_breed_age_stages").ForEach(func(breedKey, breedValue gjson.Result) bool {
if _, ok := pigTypeDescriptions[breedKey.String()]; !ok {
pigTypeDescriptions[breedKey.String()] = make(map[string]models.PigType)
}
breedValue.ForEach(func(ageStageKey, ageStageValue gjson.Result) bool {
var pt models.PigType
pt.Description = ageStageValue.Get("description").String()
pt.DailyFeedIntake = float32(ageStageValue.Get("daily_feed_intake").Float())
pt.DailyGainWeight = float32(ageStageValue.Get("daily_gain_weight").Float())
pt.MinDays = uint32(ageStageValue.Get("min_days").Uint())
pt.MaxDays = uint32(ageStageValue.Get("max_days").Uint())
pt.MinWeight = float32(ageStageValue.Get("min_weight").Float())
pt.MaxWeight = float32(ageStageValue.Get("max_weight").Float())
pigTypeDescriptions[breedKey.String()][ageStageKey.String()] = pt
return true
})
return true
})
}
// 3. 将通过校验的、干净的数据写入数据库
dataNode := gjson.GetBytes(jsonData, "data")
var err error // 用于捕获 ForEach 内部的错误
dataNode.ForEach(func(breedKey, breedValue gjson.Result) bool {
breedName := breedKey.String()
var pigBreed models.PigBreed
// 查找或创建 PigBreed
pbDesc := pigBreedDescriptions[breedName]
err = tx.Where(models.PigBreed{Name: breedName}).
FirstOrCreate(&pigBreed, models.PigBreed{
Name: breedName,
Description: pbDesc.Description,
ParentInfo: pbDesc.ParentInfo,
AppearanceFeatures: pbDesc.AppearanceFeatures,
BreedAdvantages: pbDesc.BreedAdvantages,
BreedDisadvantages: pbDesc.BreedDisadvantages,
}).Error
if err != nil {
return false
}
breedValue.ForEach(func(ageStageKey, ageStageValue gjson.Result) bool {
ageStageName := ageStageKey.String()
var pigAgeStage models.PigAgeStage
// 查找或创建 PigAgeStage
pasDesc := pigAgeStageDescriptions[ageStageName]
err = tx.Where(models.PigAgeStage{Name: ageStageName}).
FirstOrCreate(&pigAgeStage, models.PigAgeStage{
Name: ageStageName,
Description: pasDesc.Description,
}).Error
if err != nil {
return false
}
var pigType models.PigType
// 查找或创建 PigType
ptDesc := pigTypeDescriptions[breedName][ageStageName]
err = tx.Where(models.PigType{BreedID: pigBreed.ID, AgeStageID: pigAgeStage.ID}).
FirstOrCreate(&pigType, models.PigType{
BreedID: pigBreed.ID,
AgeStageID: pigAgeStage.ID,
Description: ptDesc.Description,
DailyFeedIntake: ptDesc.DailyFeedIntake,
DailyGainWeight: ptDesc.DailyGainWeight,
MinDays: ptDesc.MinDays,
MaxDays: ptDesc.MaxDays,
MinWeight: ptDesc.MinWeight,
MaxWeight: ptDesc.MaxWeight,
}).Error
if err != nil {
return false
}
ageStageValue.ForEach(func(nutrientKey, nutrientValue gjson.Result) bool {
nutrientName := nutrientKey.String()
minReq := float32(nutrientValue.Get("min_requirement").Float())
maxReq := float32(nutrientValue.Get("max_requirement").Float())
var nutrient models.Nutrient
// 查找或创建 Nutrient (这里假设 Nutrient 已经在 seedNutrients 中处理,但为了健壮性,再次 FirstOrCreate)
err = tx.Where(models.Nutrient{Name: nutrientName}).
FirstOrCreate(&nutrient, models.Nutrient{
Name: nutrientName,
// Description 字段在 nutrient seeder 中处理,这里不设置
}).Error
if err != nil {
return false
}
linkData := models.PigNutrientRequirement{
PigTypeID: pigType.ID,
NutrientID: nutrient.ID,
MinRequirement: minReq,
MaxRequirement: maxReq,
}
// 使用 FirstOrCreate 确保关联的唯一性
err = tx.Where(models.PigNutrientRequirement{
PigTypeID: pigType.ID,
NutrientID: nutrient.ID,
}).FirstOrCreate(&linkData, linkData).Error
if err != nil {
return false
}
return true
})
return err == nil // 如果内部遍历有错误,则停止外部遍历
})
return err == nil // 如果内部遍历有错误,则停止外部遍历
})
return err // 返回捕获到的错误
}
// validateAndParsePigNutrientRequirementJSON 严格校验猪营养需求JSON文件
func validateAndParsePigNutrientRequirementJSON(jsonData []byte) error {
dataNode := gjson.GetBytes(jsonData, "data")
if !dataNode.Exists() {
return errors.New("JSON文件中缺少 'data' 字段")
}
if !dataNode.IsObject() {
return errors.New("'data' 字段必须是一个JSON对象")
}
decoder := json.NewDecoder(bytes.NewReader([]byte(dataNode.Raw)))
decoder.UseNumber()
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return fmt.Errorf("'data' 字段解析起始符失败: %v", err)
}
seenBreeds := make(map[string]bool)
for decoder.More() {
// 解析 PigBreed 名称
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("解析猪品种名称失败: %w", err)
}
breedName := t.(string)
if seenBreeds[breedName] {
return fmt.Errorf("猪品种名称 '%s' 重复", breedName)
}
seenBreeds[breedName] = true
// 解析该品种的年龄阶段对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return fmt.Errorf("期望猪品种 '%s' 的值是一个JSON对象", breedName)
}
seenAgeStages := make(map[string]bool)
for decoder.More() {
// 解析 PigAgeStage 名称
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("在猪品种 '%s' 中解析年龄阶段名称失败: %w", breedName, err)
}
ageStageName := t.(string)
if seenAgeStages[ageStageName] {
return fmt.Errorf("在猪品种 '%s' 中, 年龄阶段名称 '%s' 重复", breedName, ageStageName)
}
seenAgeStages[ageStageName] = true
// 解析该年龄阶段的营养成分对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return fmt.Errorf("期望年龄阶段 '%s' 的值是一个JSON对象", ageStageName)
}
seenNutrients := make(map[string]bool)
for decoder.More() {
// 解析 Nutrient 名称
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("在年龄阶段 '%s' 中解析营养素名称失败: %w", ageStageName, err)
}
nutrientName := t.(string)
if seenNutrients[nutrientName] {
return fmt.Errorf("在年龄阶段 '%s' 中, 营养素名称 '%s' 重复", ageStageName, nutrientName)
}
seenNutrients[nutrientName] = true
// 解析 min_requirement 和 max_requirement 对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return fmt.Errorf("期望营养素 '%s' 的值是一个JSON对象", nutrientName)
}
for decoder.More() {
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("解析营养素 '%s' 的需求键失败: %w", nutrientName, err)
}
// key := t.(string) // 校验时不需要使用 key 的值
t, err = decoder.Token()
if err != nil {
return fmt.Errorf("解析营养素 '%s' 的需求值失败: %w", nutrientName, err)
}
if _, ok := t.(json.Number); !ok {
return fmt.Errorf("期望营养素 '%s' 的需求值是数字, 但实际得到的类型是 %T, 值为 '%v'", nutrientName, t, t)
}
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return fmt.Errorf("解析营养素 '%s' 的值结束符 '}' 失败", nutrientName)
}
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return fmt.Errorf("解析年龄阶段 '%s' 的值结束符 '}' 失败", ageStageName)
}
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return fmt.Errorf("解析猪品种 '%s' 的值结束符 '}' 失败", breedName)
}
}
return nil
}
// validateAndParseNutrientJSON 严格校验JSON文件
func validateAndParseNutrientJSON(jsonData []byte) error {
dataNode := gjson.GetBytes(jsonData, "data")
if !dataNode.Exists() {
return errors.New("JSON文件中缺少 'data' 字段")
}
if !dataNode.IsObject() {
return errors.New("'data' 字段必须是一个JSON对象")
}
decoder := json.NewDecoder(bytes.NewReader([]byte(dataNode.Raw)))
decoder.UseNumber()
// 读取 "{"
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return errors.New("'data' 字段解析起始符失败")
}
seenRawMaterials := make(map[string]bool)
for decoder.More() {
// 1. 解析原料名称
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("解析原料名称失败: %w", err)
}
rawMaterialName := t.(string)
if seenRawMaterials[rawMaterialName] {
return fmt.Errorf("原料名称 '%s' 重复", rawMaterialName)
}
seenRawMaterials[rawMaterialName] = true
// 2. 解析该原料的营养成分对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return fmt.Errorf("期望原料 '%s' 的值是一个JSON对象", rawMaterialName)
}
seenNutrients := make(map[string]bool)
for decoder.More() {
// 解析营养素名称
t, err := decoder.Token()
if err != nil {
return fmt.Errorf("在原料 '%s' 中解析营养素名称失败: %w", rawMaterialName, err)
}
nutrientName := t.(string)
if seenNutrients[nutrientName] {
return fmt.Errorf("在原料 '%s' 中, 营养素名称 '%s' 重复", rawMaterialName, nutrientName)
}
seenNutrients[nutrientName] = true
// 解析营养素含量
t, err = decoder.Token()
if err != nil {
return fmt.Errorf("在原料 '%s' 中解析营养素 '%s' 的含量值失败: %w", rawMaterialName, nutrientName, err)
}
if _, ok := t.(json.Number); !ok {
return fmt.Errorf("期望营养素 '%s' 的含量值是数字, 但实际得到的类型是 %T, 值为 '%v'", nutrientName, t, t)
}
}
// 读取营养成分对象的 "}"
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return fmt.Errorf("解析原料 '%s' 的值结束符 '}' 失败", rawMaterialName)
}
}
return nil
}