127 lines
4.2 KiB
Go
127 lines
4.2 KiB
Go
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 // 提交事务
|
||
})
|
||
}
|