146 lines
5.3 KiB
Go
146 lines
5.3 KiB
Go
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
|
||
})
|
||
}
|