gohub/pkg/migrate/migrator.go

// Package migrate  操作对象
// 负责创建 migrations 数据表，以及执行迁移操作。
package migrate

import (
	"os"

	"gorm.io/gorm"

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

// Migrator 结构体用于存储迁移器的相关信息。
type Migrator struct {
	Folder   string        // Folder 是存储迁移文件的目录名。
	DB       *gorm.DB      // Db 是数据库连接对象。
	Migrator gorm.Migrator // Migrator 是 GORM 的迁移器对象。
}

func (m *Migrator) createMigrationsTable() {
	migration := Migration{}

	if !m.Migrator.HasTable(&migration) {
		// 如果表不存在，则创建表
		if err := m.DB.Set("gorm:table_options", "ENGINE=InnoDB CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci").Migrator().CreateTable(&migration); err != nil {
			console.ExitIf(err)
			return
		}
	}
}

// Migration 对应数据表 migrations 里的一条记录
type Migration struct {
	ID        uint64 `gorm:"primaryKey;autoIncrement;"`
	Migration string `gorm:"type:varchar(255);not null;unique;"`
	Batch     int
}

func NewMigrator() *Migrator {
	migrator := &Migrator{
		Folder:   "database/migrations",
		DB:       database.DB,
		Migrator: database.DB.Migrator(),
	}

	migrator.createMigrationsTable()

	return migrator
}

// Up 执行所有未执行的迁移文件。
func (m *Migrator) Up() {

	// 读取所有迁移文件， 确保按照时间排序
	migrateFiles := m.readAllMigrationFiles()
	batch := m.getBatch()

	// 获取所有迁移数据
	var migrations []Migration
	m.DB.Find(&migrations)

	// 可以通过此值来判断数据库是否已是最新
	runed := false

	// 遍历迁移文件，执行未执行的迁移文件
	for _, mfile := range migrateFiles {

		if mfile.isNotMigrated(migrations) {
			m.runUpMigration(mfile, batch)
			runed = true
		}
	}

	if !runed {
		console.Success("database is up to date")
	}
}

// Rollback 回滚上一次的迁移操作
func (m *Migrator) Rollback() {

	// 获取最后一批次的迁移数据
	lastMigration := Migration{}
	m.DB.Order("id desc").First(&lastMigration)
	var migrations []Migration
	m.DB.Where("batch = ?", lastMigration.Batch).Order("id desc").Find(&migrations)

	// 回滚迁移操作
	if !m.rollbackMigrations(migrations) {
		console.Success("[migrations] table is empty, nothing to rollback")
	}
}

// Reset 回滚所有的迁移操作
func (m *Migrator) Reset() {

	// 获取所有的迁移数据
	var migrations []Migration
	m.DB.Order("id desc").Find(&migrations)

	// 回滚迁移操作
	if !m.rollbackMigrations(migrations) {
		console.Success("[migrations] table is empty, nothing to rollback")
	}
}

// Refresh 回滚所有的迁移操作， 并重新执行所有的迁移操作
func (m *Migrator) Refresh() {
	// 回滚所有的迁移操作
	m.Reset()
	// 重新执行所有的迁移操作
	m.Up()
}

// Fresh Drop 所有的表， 并重新执行所有的迁移操作
func (m *Migrator) Fresh() {

	// 获取数据库名称，用以提示
	dbname := database.CurrentDatabase()

	// 删除所有表
	err := database.DeleteAllTables()
	console.ExitIf(err)
	console.Success("database " + dbname + " cleared")

	// 重新创建 migrates 表
	m.createMigrationsTable()
	console.Success("migrations table created")

	// 重新执行所有的迁移操作
	m.Up()
}

// 回滚迁移操作
func (m *Migrator) rollbackMigrations(migrations []Migration) bool {

	// 标记是否真的有执行了迁移回退的操作
	runed := false
	// 遍历迁移数据，回滚迁移操作
	for _, migration := range migrations {
		// 友好提示
		console.Warning("rolling back " + migration.Migration + " ...")
		// 获取迁移文件
		mfile := getMigrationFile(migration.Migration)
		if mfile.Down != nil {
			// 执行迁移回退操作
			mfile.Down(database.DB.Migrator(), database.DB)
		}
		runed = true

		// 删除迁移数据
		m.DB.Delete(&migration)
		console.Success("rolled back " + migration.Migration + " finished")
	}

	return runed
}

// 获取当前这个批次的值
func (m *Migrator) getBatch() int {

	batch := 1

	lastMigration := Migration{}
	m.DB.Order("id desc").First(&lastMigration)

	// 如果有值的话，加一
	if lastMigration.ID > 0 {
		batch = lastMigration.Batch + 1
	}

	return batch
}

// 从文件目录读取文件，保证正确的时间排序
func (m *Migrator) readAllMigrationFiles() []MigrationFile {
	// 读取 database/migrations 目录下的所有迁移文件
	files, err := os.ReadDir(m.Folder)
	console.ExitIf(err)

	var migrateFiles []MigrationFile
	// 遍历所有迁移文件
	for _, f := range files {

		// 去除文件后缀 .go
		fileName := file.NameWithoutExtension(f.Name())
		// 通过迁移文件名称获取[MigrationFile]结构体
		mfile := getMigrationFile(fileName)
		if len(mfile.FileName) > 0 {
			migrateFiles = append(migrateFiles, mfile)
		}
	}
	return migrateFiles
}

// 插入一条记录到数据库，同时执行迁移操作
// 此方法会在 Up 方法中被调用
// mfile 是迁移文件的结构体，batch 是批次号，用于区分不同的迁移批次。
// 此方法会执行迁移操作，同时插入一条记录到数据库。
// 如果迁移操作成功，会输出一条成功信息；如果迁移操作失败，会输出一条错误信息，并退出程序。
func (m *Migrator) runUpMigration(file MigrationFile, batch int) {

	if file.Up != nil {
		console.Warning("migrating " + file.FileName)
		// 执行迁移操作
		file.Up(database.DB.Migrator(), database.DB)
		console.Success("migrated " + file.FileName)
	}
	// 插入一条记录到数据库
	err := m.DB.Create(&Migration{
		Migration: file.FileName,
		Batch:     batch,
	}).Error
	console.ExitIf(err)
}