2024-03-31 20:09:20 +08:00

75 lines
1.6 KiB
Go

package model
import (
"github.com/loveuer/nfflow/internal/database"
"github.com/loveuer/nfflow/internal/opt"
"github.com/loveuer/nfflow/internal/util"
"github.com/sirupsen/logrus"
"strings"
)
func MustInit(db database.Store) {
var err error
if err = initModel(db); err != nil {
logrus.Fatalf("model.MustInit: init models err=%v", err)
}
logrus.Info("MustInitModels: auto_migrate privilege model success")
if err = initData(db); err != nil {
logrus.Fatalf("model.MustInit: init datas err=%v", err)
}
}
func initModel(client database.Store) error {
if err := client.Session(util.Timeout(10)).AutoMigrate(
&User{},
&OpLog{},
&Task{},
); err != nil {
return err
}
logrus.Info("MustInitModels: auto_migrate user model success")
return nil
}
func initData(client database.Store) error {
var (
err error
)
{
count := 0
if err = client.Session(util.Timeout()).Model(&User{}).Select("count(id)").Take(&count).Error; err != nil {
return err
}
if count < len(initUsers) {
logrus.Warn("mustInitDatas: user count = 0, start init...")
for _, user := range initUsers {
if err = client.Session(util.Timeout(5)).Model(&User{}).Create(user).Error; err != nil {
if !strings.Contains(err.Error(), "SQLSTATE 23505") {
return err
}
}
}
if opt.Cfg.Database.Type == "postgresql" {
if err = client.Session(util.Timeout(5)).Exec(`SELECT setval('users_id_seq', (SELECT MAX(id) FROM users))`).Error; err != nil {
return err
}
}
logrus.Info("mustInitDatas: creat init users success")
}
}
return nil
}