xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-07 17:23:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/upstream/main'

Browse Source
This commit is contained in:
Laisky.Cai
2024-01-28 13:44:16 +00:00
parent 50e0638313 2cd1a78203
commit 10cbe63c0d
113 changed files with 3608 additions and 3013 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
model/main.go
View File

@@ -2,24 +2,28 @@ package model
import (
"fmt"
"github.com/Laisky/errors/v2"
"os"
"strings"
"time"
"github.com/pkg/errors"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"one-api/common"
"os"
"strings"
"time"
)
var DB *gorm.DB
func createRootAccountIfNeed() error {
var user User
//if user.Status != common.UserStatusEnabled {
//if user.Status != util.UserStatusEnabled {
if err := DB.First(&user).Error; err != nil {
common.SysLog("no user exists, create a root user for you: username is root, password is 123456")
logger.SysLog("no user exists, create a root user for you: username is root, password is 123456")
hashedPassword, err := common.Password2Hash("123456")
if err != nil {
return errors.WithStack(err)
@@ -30,7 +34,7 @@ func createRootAccountIfNeed() error {
Role: common.RoleRootUser,
Status: common.UserStatusEnabled,
DisplayName: "Root User",
AccessToken: common.GetUUID(),
AccessToken: helper.GetUUID(),
Quota: 100000000,
}
DB.Create(&rootUser)
@@ -43,7 +47,7 @@ func chooseDB() (*gorm.DB, error) {
dsn := os.Getenv("SQL_DSN")
if strings.HasPrefix(dsn, "postgres://") {
// Use PostgreSQL
common.SysLog("using PostgreSQL as database")
logger.SysLog("using PostgreSQL as database")
common.UsingPostgreSQL = true
return gorm.Open(postgres.New(postgres.Config{
DSN: dsn,
@@ -53,13 +57,13 @@ func chooseDB() (*gorm.DB, error) {
})
}
// Use MySQL
common.SysLog("using MySQL as database")
logger.SysLog("using MySQL as database")
return gorm.Open(mysql.Open(dsn), &gorm.Config{
PrepareStmt: true, // precompile SQL
})
}
// Use SQLite
common.SysLog("SQL_DSN not set, using SQLite as database")
logger.SysLog("SQL_DSN not set, using SQLite as database")
common.UsingSQLite = true
config := fmt.Sprintf("?_busy_timeout=%d", common.SQLiteBusyTimeout)
return gorm.Open(sqlite.Open(common.SQLitePath+config), &gorm.Config{
@@ -70,7 +74,7 @@ func chooseDB() (*gorm.DB, error) {
func InitDB() (err error) {
db, err := chooseDB()
if err == nil {
if common.DebugEnabled {
if config.DebugEnabled {
db = db.Debug()
}
DB = db
@@ -78,14 +82,14 @@ func InitDB() (err error) {
if err != nil {
return errors.WithStack(err)
}
sqlDB.SetMaxIdleConns(common.GetOrDefault("SQL_MAX_IDLE_CONNS", 100))
sqlDB.SetMaxOpenConns(common.GetOrDefault("SQL_MAX_OPEN_CONNS", 1000))
sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetOrDefault("SQL_MAX_LIFETIME", 60)))
sqlDB.SetMaxIdleConns(helper.GetOrDefaultEnvInt("SQL_MAX_IDLE_CONNS", 100))
sqlDB.SetMaxOpenConns(helper.GetOrDefaultEnvInt("SQL_MAX_OPEN_CONNS", 1000))
sqlDB.SetConnMaxLifetime(time.Second * time.Duration(helper.GetOrDefaultEnvInt("SQL_MAX_LIFETIME", 60)))
if !common.IsMasterNode {
if !config.IsMasterNode {
return nil
}
common.SysLog("database migration started")
logger.SysLog("database migration started")
err = db.AutoMigrate(&Channel{})
if err != nil {
return errors.WithStack(err)
@@ -114,11 +118,11 @@ func InitDB() (err error) {
if err != nil {
return errors.WithStack(err)
}
common.SysLog("database migrated")
logger.SysLog("database migrated")
err = createRootAccountIfNeed()
return errors.WithStack(err)
} else {
common.FatalLog(err)
logger.FatalLog(err)
}
return errors.WithStack(err)
}