device-manage-server/device-manage/cmd/migrate/migration/init.go

package migration

import (
	"fmt"
	"log"
	"path/filepath"
	"sort"
	"strconv"
	"sync"

	"github.com/spf13/cast"
	"gorm.io/gorm"
)

var Migrate = &Migration{
	version: make(map[int]func(db *gorm.DB, version string) error),
}

type Migration struct {
	db      *gorm.DB
	version map[int]func(db *gorm.DB, version string) error
	mutex   sync.Mutex
}

func (e *Migration) GetDb() *gorm.DB {
	return e.db
}

func (e *Migration) SetDb(db *gorm.DB) {
	e.db = db
}

//version 初始化 init 时调用, 将数据库信息导入到 migration.
func (e *Migration) SetVersion(k int, f func(db *gorm.DB, version string) error) {
	e.mutex.Lock()
	defer e.mutex.Unlock()
	e.version[k] = f
}

func (e *Migration) Migrate() {
	versions := make([]int, 0)
	for k := range e.version {
		versions = append(versions, k)
	}
	sort.IntsAreSorted(versions)
	var err error
	var count int64
	for _, v := range versions {
		fmt.Println(v)
		err = e.db.Debug().Table("sys_migration").Where("version = ?", v).Count(&count).Error
		if err != nil {
			log.Fatalln(err)
		}
		if count > 0 {
			log.Println(count)
			count = 0
			continue
		}
		err = (e.version[v])(e.db.Debug(), strconv.Itoa(v))
		if err != nil {
			log.Fatalln(err)
		}
	}
}

func GetFilename(s string) int {
	s = filepath.Base(s)
	fmt.Println(s)
	return cast.ToInt(s[:13])
}