232 lines
7.5 KiB
Go
232 lines
7.5 KiB
Go
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")
|
||
|
||
// 定义必须存在的预设数据类型
|
||
requiredTypes := []string{"nutrient"}
|
||
processedTypes := make(map[string]bool)
|
||
|
||
// 用于检测重复的 type
|
||
typeToFileMap := make(map[string]string)
|
||
|
||
files, err := os.ReadDir(presetDir)
|
||
if err != nil {
|
||
return fmt.Errorf("读取预设数据目录 '%s' 失败: %w", presetDir, err)
|
||
}
|
||
|
||
return db.Transaction(func(tx *gorm.DB) error {
|
||
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()
|
||
|
||
if existingFile, found := typeToFileMap[dataTypeStr]; found {
|
||
return fmt.Errorf("预设数据校验失败: type '%s' 在文件 '%s' 和 '%s' 中重复定义", dataTypeStr, existingFile, filePath)
|
||
}
|
||
typeToFileMap[dataTypeStr] = filePath
|
||
|
||
var seederFunc SeederFunc
|
||
switch dataTypeStr {
|
||
case "nutrient":
|
||
seederFunc = seedNutrients
|
||
default:
|
||
logger.Warnf("警告: 文件 '%s' 中存在未知的 type: '%s'，已跳过", filePath, dataTypeStr)
|
||
continue
|
||
}
|
||
|
||
if err := seederFunc(tx, jsonData); err != nil {
|
||
return fmt.Errorf("处理文件 '%s' (type: %s) 时发生错误: %w", filePath, dataTypeStr, 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
|
||
}
|
||
|
||
// 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
|
||
}
|