package database import ( "context" "fmt" "os" "path/filepath" "strings" "git.huangwc.com/pig/pig-farm-controller/internal/infra/logs" "github.com/tidwall/gjson" "gorm.io/gorm" ) // Seeder 定义了数据播种器的通用接口。 // 每个播种器负责处理一种特定类型的预设数据。 type Seeder interface { // Type 返回此播种器能处理的唯一类型标识符, // 这个标识符应与预设 JSON 文件中的 "type" 字段相匹配。 Type() string // IsRequired 标记此播种器对应的预设文件是否为必需的。 // 如果为 true 且对应的 .json 文件不存在,整个播种过程将失败并报错。 IsRequired() bool // Seed 执行具体的数据播种逻辑。 // 它接收数据库事务、上下文和从 JSON 文件读取的原始字节数据。 Seed(ctx context.Context, tx *gorm.DB, jsonData []byte) error } // SeedFromPreset 是一个通用的数据播种函数。 // 它会读取指定目录下的所有 .json 文件,并根据文件内容中的 "type" 字段, // 将其分发给由外部注入的相应 Seeder 进行处理。 // Seeder 的处理顺序由其在注入切片中的顺序决定。 func SeedFromPreset(ctx context.Context, db *gorm.DB, presetDir string, seeders []Seeder) error { seedCtx, logger := logs.Trace(ctx, ctx, "SeedFromPreset") // --- 步骤 1: 校验注入的 Seeders 是否有重复类型 --- tempSeederTypes := make(map[string]bool) for _, s := range seeders { if tempSeederTypes[s.Type()] { return fmt.Errorf("播种器初始化失败: 存在重复的 Seeder 类型 '%s'", s.Type()) } tempSeederTypes[s.Type()] = true } // --- 步骤 2: 读取并分组所有预设文件 --- groupedFiles := make(map[string][][]byte) typeToFileMap := make(map[string]string) files, err := os.ReadDir(presetDir) if err != nil { if os.IsNotExist(err) { // 目录不存在 // 检查是否有必需的 Seeder,如果目录不存在但需要文件,则报错 var requiredSeederTypes []string for _, s := range seeders { if s.IsRequired() { requiredSeederTypes = append(requiredSeederTypes, s.Type()) } } if len(requiredSeederTypes) > 0 { return fmt.Errorf("预设数据校验失败: 预设目录 '%s' 不存在, 但系统需要以下必需类型: [%s]", presetDir, strings.Join(requiredSeederTypes, ", ")) } logger.Warnf("预设数据目录 '%s' 不存在,跳过播种。", presetDir) return nil } return fmt.Errorf("读取预设数据目录 '%s' 失败: %w", presetDir, err) } // 处理目录为空的情况 if len(files) == 0 { var requiredSeederTypes []string for _, s := range seeders { if s.IsRequired() { requiredSeederTypes = append(requiredSeederTypes, s.Type()) } } if len(requiredSeederTypes) > 0 { return fmt.Errorf("预设数据校验失败: 预设目录 '%s' 为空, 但系统需要以下必需类型: [%s]", presetDir, strings.Join(requiredSeederTypes, ", ")) } logger.Warnf("预设数据目录 '%s' 为空,跳过播种。", presetDir) return nil } // 读取所有文件并按 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) } // --- 步骤 3: 在事务中按注入顺序执行播种,并采用快速失败策略 --- return db.Transaction(func(tx *gorm.DB) error { logger.Info("开始执行数据播种事务...") // 直接遍历注入的 seeders,顺序由调用方保证 for _, seeder := range seeders { dataTypeStr := seeder.Type() jsonDatas, fileExists := groupedFiles[dataTypeStr] if fileExists { // --- 文件存在,执行播种 --- logger.Infof("正在使用播种器处理类型: '%s'...", dataTypeStr) for _, jsonData := range jsonDatas { originalFilePath := typeToFileMap[dataTypeStr] if err := seeder.Seed(seedCtx, tx, jsonData); err != nil { return fmt.Errorf("处理文件 (type: %s, path: %s) 时发生错误: %w", dataTypeStr, originalFilePath, err) } } logger.Infof("类型 '%s' 处理完成。", dataTypeStr) } else { // --- 文件不存在,检查是否必需 --- if seeder.IsRequired() { // 快速失败:如果是必需的但文件不存在,立即报错并回滚事务 return fmt.Errorf("预设数据校验失败: 缺少必需的预设文件类型: '%s'", dataTypeStr) } // 如果不是必需的,则只记录日志 logger.Infof("未找到可选类型为 '%s' 的预设文件,跳过该播种器。", dataTypeStr) } } logger.Info("数据播种事务成功完成。") return nil }) }