gohub/pkg/seed/seeder.go

// Package seed 存放数据库填充数据
package seed

import (
	"gorm.io/gorm"

	"github.com/runningwater/gohub/pkg/console"
	"github.com/runningwater/gohub/pkg/database"
)

// 存放所有 Seeder
var seeders []Seeder

// 按排序执行的 Seeder 数组
var orderedSeederNames []string

type SeederFunc func(*gorm.DB)

// Seeder 对应每一个 database/seeders 目录下的 Seeder 文件
type Seeder struct {
	Name string
	Func SeederFunc
}

// Add AddSeeder 添加一个 Seeder
func Add(name string, fn SeederFunc) {
	seeders = append(seeders, Seeder{
		Name: name,
		Func: fn,
	})
}

// SetRunOrder 设置 按顺序执行的 Seeder 数组
func SetRunOrder(names []string) {
	orderedSeederNames = names
}

// GetSeeder 获取 Seeder
func GetSeeder(name string) Seeder {
	for _, seeder := range seeders {
		if seeder.Name == name {
			return seeder
		}
	}
	return Seeder{}
}

// RunAll 执行所有 Seeder
func RunAll() {
	// 先运行 ordered 的
	executed := make(map[string]string)
	for _, name := range orderedSeederNames {
		seeder := GetSeeder(name)
		if len(seeder.Name) > 0 {
			if seeder.Func != nil {
				seeder.Func(database.DB)
				executed[name] = name
			}
		}
	}

	// 再运行剩下的
	for _, seeder := range seeders {
		if _, ok := executed[seeder.Name]; !ok {
			if seeder.Func != nil {
				console.Warning("Running Seeder: " + seeder.Name)
				seeder.Func(database.DB)
			}
		}
	}
}

// RunSeeder 运行指定的 Seeder
func RunSeeder(name string) {
	seeder := GetSeeder(name)
	if seeder.Func != nil {
		console.Warning("Running Seeder: " + seeder.Name)
		seeder.Func(database.DB)
	}
}