package database import ( "context" "fmt" "os" "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 } // SeedFromPreset 是一个通用的数据播种函数。 // 它会读取指定目录下的所有 .json 文件,并根据文件内容中的 "type" 字段进行分发。 // 同时,它会校验所有必需的预设类型是否都已成功加载。 func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string) 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 } 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 = 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) } } if len(missingTypes) > 0 { return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: [%s]", strings.Join(missingTypes, ", ")) } return nil // 提交事务 }) }