Alireza Ahmadi
86 lines
1.2 KiB
Go
86 lines
1.2 KiB
Go
package database
|
|
|
|
import (
|
|
"os"
|
|
"path"
|
|
"s-ui/config"
|
|
"s-ui/database/model"
|
|
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
var db *gorm.DB
|
|
|
|
func initUser() error {
|
|
var count int64
|
|
err := db.Model(&model.User{}).Count(&count).Error
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if count == 0 {
|
|
user := &model.User{
|
|
Username: "admin",
|
|
Password: "admin",
|
|
}
|
|
return db.Create(user).Error
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func OpenDB(dbPath string) error {
|
|
dir := path.Dir(dbPath)
|
|
err := os.MkdirAll(dir, 01740)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
var gormLogger logger.Interface
|
|
|
|
if config.IsDebug() {
|
|
gormLogger = logger.Default
|
|
} else {
|
|
gormLogger = logger.Discard
|
|
}
|
|
|
|
c := &gorm.Config{
|
|
Logger: gormLogger,
|
|
}
|
|
db, err = gorm.Open(sqlite.Open(dbPath), c)
|
|
return err
|
|
}
|
|
|
|
func InitDB(dbPath string) error {
|
|
err := OpenDB(dbPath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = db.AutoMigrate(
|
|
&model.Setting{},
|
|
&model.Tls{},
|
|
&model.User{},
|
|
&model.Stats{},
|
|
&model.Client{},
|
|
&model.Changes{},
|
|
)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = initUser()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func GetDB() *gorm.DB {
|
|
return db
|
|
}
|
|
|
|
func IsNotFound(err error) bool {
|
|
return err == gorm.ErrRecordNotFound
|
|
}
|