pig-farm-controller/internal/infra/database/seeder.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
	})
}