Files
pig-farm-controller/internal/infra/database/seeder.go

533 lines
19 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文件检查内部重复键
parsedData, err := validateAndParseNutrientJSON(jsonData)
if 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. 将通过校验的、干净的数据写入数据库
for rawMaterialName, nutrients := range parsedData {
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 {
return fmt.Errorf("预设原料 '%s' 失败: %w", rawMaterialName, err)
}
for nutrientName, value := range nutrients {
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 {
return fmt.Errorf("预设营养素 '%s' 失败: %w", nutrientName, err)
}
linkData := models.RawMaterialNutrient{
RawMaterialID: rawMaterial.ID,
NutrientID: nutrient.ID,
}
// 使用 FirstOrCreate 确保关联的唯一性
if err := tx.Where(linkData).FirstOrCreate(&linkData, models.RawMaterialNutrient{
RawMaterialID: linkData.RawMaterialID,
NutrientID: linkData.NutrientID,
Value: value,
}).Error; err != nil {
return fmt.Errorf("为原料 '%s' 和营养素 '%s' 创建关联失败: %w", rawMaterialName, nutrientName, err)
}
}
}
return nil
}
// seedPigNutrientRequirements 先严格校验JSON源文件然后以“有则跳过”的模式播种数据。
func seedPigNutrientRequirements(tx *gorm.DB, jsonData []byte) error {
// 1. 严格校验JSON文件检查内部重复键
parsedData, err := validateAndParsePigNutrientRequirementJSON(jsonData)
if 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. 将通过校验的、干净的数据写入数据库
for breedName, ageStagesData := range parsedData {
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 fmt.Errorf("预设猪品种 '%s' 失败: %w", breedName, err)
}
for ageStageName, nutrientsData := range ageStagesData {
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 fmt.Errorf("预设猪年龄阶段 '%s' 失败: %w", ageStageName, err)
}
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 fmt.Errorf("预设猪类型 '%s' - '%s' 失败: %w", breedName, ageStageName, err)
}
for nutrientName, requirement := range nutrientsData {
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 fmt.Errorf("预设营养素 '%s' 失败: %w", nutrientName, err)
}
linkData := models.PigNutrientRequirement{
PigTypeID: pigType.ID,
NutrientID: nutrient.ID,
MinRequirement: requirement.MinRequirement,
MaxRequirement: requirement.MaxRequirement,
}
// 使用 FirstOrCreate 确保关联的唯一性
if err := tx.Where(models.PigNutrientRequirement{
PigTypeID: pigType.ID,
NutrientID: nutrient.ID,
}).FirstOrCreate(&linkData, linkData).Error; err != nil {
return fmt.Errorf("为猪类型 '%s' - '%s' 和营养素 '%s' 创建营养需求失败: %w", breedName, ageStageName, nutrientName, err)
}
}
}
}
return nil
}
// validateAndParsePigNutrientRequirementJSON 严格校验并解析猪营养需求JSON文件
func validateAndParsePigNutrientRequirementJSON(jsonData []byte) (map[string]map[string]map[string]struct {
MinRequirement float32
MaxRequirement float32
}, error) {
dataNode := gjson.GetBytes(jsonData, "data")
if !dataNode.Exists() {
return nil, errors.New("JSON文件中缺少 'data' 字段")
}
if !dataNode.IsObject() {
return nil, 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 nil, fmt.Errorf("'data' 字段解析起始符失败: %v", err)
}
result := make(map[string]map[string]map[string]struct {
MinRequirement float32
MaxRequirement float32
})
seenBreeds := make(map[string]bool)
for decoder.More() {
// 解析 PigBreed 名称
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("解析猪品种名称失败: %w", err)
}
breedName := t.(string)
if seenBreeds[breedName] {
return nil, fmt.Errorf("猪品种名称 '%s' 重复", breedName)
}
seenBreeds[breedName] = true
// 解析该品种的年龄阶段对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return nil, fmt.Errorf("期望猪品种 '%s' 的值是一个JSON对象", breedName)
}
ageStages := make(map[string]map[string]struct {
MinRequirement float32
MaxRequirement float32
})
seenAgeStages := make(map[string]bool)
for decoder.More() {
// 解析 PigAgeStage 名称
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("在猪品种 '%s' 中解析年龄阶段名称失败: %w", breedName, err)
}
ageStageName := t.(string)
if seenAgeStages[ageStageName] {
return nil, fmt.Errorf("在猪品种 '%s' 中, 年龄阶段名称 '%s' 重复", breedName, ageStageName)
}
seenAgeStages[ageStageName] = true
// 解析该年龄阶段的营养成分对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return nil, fmt.Errorf("期望年龄阶段 '%s' 的值是一个JSON对象", ageStageName)
}
nutrients := make(map[string]struct {
MinRequirement float32
MaxRequirement float32
})
seenNutrients := make(map[string]bool)
for decoder.More() {
// 解析 Nutrient 名称
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("在年龄阶段 '%s' 中解析营养素名称失败: %w", ageStageName, err)
}
nutrientName := t.(string)
if seenNutrients[nutrientName] {
return nil, fmt.Errorf("在年龄阶段 '%s' 中, 营养素名称 '%s' 重复", ageStageName, nutrientName)
}
seenNutrients[nutrientName] = true
// 解析 min_requirement 和 max_requirement 对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return nil, fmt.Errorf("期望营养素 '%s' 的值是一个JSON对象", nutrientName)
}
var req struct {
MinRequirement float32
MaxRequirement float32
}
for decoder.More() {
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("解析营养素 '%s' 的需求键失败: %w", nutrientName, err)
}
key := t.(string)
t, err = decoder.Token()
if err != nil {
return nil, fmt.Errorf("解析营养素 '%s' 的需求值失败: %w", nutrientName, err)
}
if value, ok := t.(json.Number); ok {
f64, _ := value.Float64()
if key == "min_requirement" {
req.MinRequirement = float32(f64)
} else if key == "max_requirement" {
req.MaxRequirement = float32(f64)
} else {
return nil, fmt.Errorf("营养素 '%s' 中存在未知键 '%s'", nutrientName, key)
}
} else {
return nil, fmt.Errorf("期望营养素 '%s' 的 '%s' 值是数字, 但实际得到的类型是 %T, 值为 '%v'", nutrientName, key, t, t)
}
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return nil, fmt.Errorf("解析营养素 '%s' 的值结束符 '}' 失败", nutrientName)
}
nutrients[nutrientName] = req
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return nil, fmt.Errorf("解析年龄阶段 '%s' 的值结束符 '}' 失败", ageStageName)
}
ageStages[ageStageName] = nutrients
}
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return nil, fmt.Errorf("解析猪品种 '%s' 的值结束符 '}' 失败", breedName)
}
result[breedName] = ageStages
}
return result, nil
}
// validateAndParseNutrientJSON 使用 json.Decoder 手动解析,以捕获重复的键。
func validateAndParseNutrientJSON(jsonData []byte) (map[string]map[string]float32, error) {
dataNode := gjson.GetBytes(jsonData, "data")
if !dataNode.Exists() {
return nil, errors.New("JSON文件中缺少 'data' 字段")
}
if !dataNode.IsObject() {
return nil, 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 nil, errors.New("'data' 字段解析起始符失败")
}
result := make(map[string]map[string]float32)
seenRawMaterials := make(map[string]bool)
for decoder.More() {
// 1. 解析原料名称
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("解析原料名称失败: %w", err)
}
rawMaterialName := t.(string)
if seenRawMaterials[rawMaterialName] {
return nil, fmt.Errorf("原料名称 '%s' 重复", rawMaterialName)
}
seenRawMaterials[rawMaterialName] = true
// 2. 解析该原料的营养成分对象
if t, err := decoder.Token(); err != nil || t != json.Delim('{') {
return nil, fmt.Errorf("期望原料 '%s' 的值是一个JSON对象", rawMaterialName)
}
nutrients := make(map[string]float32)
seenNutrients := make(map[string]bool)
for decoder.More() {
// 解析营养素名称
t, err := decoder.Token()
if err != nil {
return nil, fmt.Errorf("在原料 '%s' 中解析营养素名称失败: %w", rawMaterialName, err)
}
nutrientName := t.(string)
if seenNutrients[nutrientName] {
return nil, fmt.Errorf("在原料 '%s' 中, 营养素名称 '%s' 重复", rawMaterialName, nutrientName)
}
seenNutrients[nutrientName] = true
// 解析营养素含量
t, err = decoder.Token()
if err != nil {
return nil, fmt.Errorf("在原料 '%s' 中解析营养素 '%s' 的含量值失败: %w", rawMaterialName, nutrientName, err)
}
if value, ok := t.(json.Number); ok {
f64, _ := value.Float64()
nutrients[nutrientName] = float32(f64)
} else {
return nil, fmt.Errorf("期望营养素 '%s' 的含量值是数字, 但实际得到的类型是 %T, 值为 '%v'", nutrientName, t, t)
}
}
// 读取营养成分对象的 "}"
if t, err := decoder.Token(); err != nil || t != json.Delim('}') {
return nil, fmt.Errorf("解析原料 '%s' 的值结束符 '}' 失败", rawMaterialName)
}
result[rawMaterialName] = nutrients
}
return result, nil
}