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' 字段,已跳过\n", 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',已跳过\n", 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. 将通过校验的、干净的数据写入数据库 for rawMaterialName, nutrients := range parsedData { var rawMaterial models.RawMaterial if err := tx.Where(models.RawMaterial{Name: rawMaterialName}).FirstOrCreate(&rawMaterial).Error; err != nil { return fmt.Errorf("预设原料 '%s' 失败: %w", rawMaterialName, err) } for nutrientName, value := range nutrients { var nutrient models.Nutrient if err := tx.Where(models.Nutrient{Name: nutrientName}).FirstOrCreate(&nutrient).Error; 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 }