Compare commits

...

3 Commits

Author SHA1 Message Date
Leo Q
ec6ad24810 feat: support smtp without auth (#1101) 2024-07-03 22:23:49 +08:00
LinZeliang
c4fe57c165 feat: support one or more log file (#1400)
Co-authored-by: Laisky.Cai <github@laisky.com>
2024-07-03 20:53:29 +08:00
igophper
274fcf3d76 refactor: init db (#1590)
Co-authored-by: 江杭辉 <jianghanghui@k.app>
2024-07-03 20:50:40 +08:00
5 changed files with 181 additions and 99 deletions

View File

@@ -145,6 +145,9 @@ var InitialRootToken = os.Getenv("INITIAL_ROOT_TOKEN")
var GeminiVersion = env.String("GEMINI_VERSION", "v1") var GeminiVersion = env.String("GEMINI_VERSION", "v1")
var OnlyOneLogFile = env.Bool("ONLY_ONE_LOG_FILE", false)
var RelayProxy = env.String("RELAY_PROXY", "") var RelayProxy = env.String("RELAY_PROXY", "")
var UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "") var UserContentRequestProxy = env.String("USER_CONTENT_REQUEST_PROXY", "")
var UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30) var UserContentRequestTimeout = env.Int("USER_CONTENT_REQUEST_TIMEOUT", 30)

View File

@@ -27,7 +27,12 @@ var setupLogOnce sync.Once
func SetupLogger() { func SetupLogger() {
setupLogOnce.Do(func() { setupLogOnce.Do(func() {
if LogDir != "" { if LogDir != "" {
logPath := filepath.Join(LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102"))) var logPath string
if config.OnlyOneLogFile {
logPath = filepath.Join(LogDir, "oneapi.log")
} else {
logPath = filepath.Join(LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102")))
}
fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil { if err != nil {
log.Fatal("failed to open log file") log.Fatal("failed to open log file")

View File

@@ -6,11 +6,16 @@ import (
"encoding/base64" "encoding/base64"
"fmt" "fmt"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"net"
"net/smtp" "net/smtp"
"strings" "strings"
"time" "time"
) )
func shouldAuth() bool {
return config.SMTPAccount != "" || config.SMTPToken != ""
}
func SendEmail(subject string, receiver string, content string) error { func SendEmail(subject string, receiver string, content string) error {
if receiver == "" { if receiver == "" {
return fmt.Errorf("receiver is empty") return fmt.Errorf("receiver is empty")
@@ -41,16 +46,24 @@ func SendEmail(subject string, receiver string, content string) error {
"Date: %s\r\n"+ "Date: %s\r\n"+
"Content-Type: text/html; charset=UTF-8\r\n\r\n%s\r\n", "Content-Type: text/html; charset=UTF-8\r\n\r\n%s\r\n",
receiver, config.SystemName, config.SMTPFrom, encodedSubject, messageId, time.Now().Format(time.RFC1123Z), content)) receiver, config.SystemName, config.SMTPFrom, encodedSubject, messageId, time.Now().Format(time.RFC1123Z), content))
auth := smtp.PlainAuth("", config.SMTPAccount, config.SMTPToken, config.SMTPServer) auth := smtp.PlainAuth("", config.SMTPAccount, config.SMTPToken, config.SMTPServer)
addr := fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort) addr := fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort)
to := strings.Split(receiver, ";") to := strings.Split(receiver, ";")
if config.SMTPPort == 465 { if config.SMTPPort == 465 || !shouldAuth() {
tlsConfig := &tls.Config{ // need advanced client
InsecureSkipVerify: true, var conn net.Conn
ServerName: config.SMTPServer, var err error
if config.SMTPPort == 465 {
tlsConfig := &tls.Config{
InsecureSkipVerify: true,
ServerName: config.SMTPServer,
}
conn, err = tls.Dial("tcp", fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort), tlsConfig)
} else {
conn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort))
} }
conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort), tlsConfig)
if err != nil { if err != nil {
return err return err
} }
@@ -59,8 +72,10 @@ func SendEmail(subject string, receiver string, content string) error {
return err return err
} }
defer client.Close() defer client.Close()
if err = client.Auth(auth); err != nil { if shouldAuth() {
return err if err = client.Auth(auth); err != nil {
return err
}
} }
if err = client.Mail(config.SMTPFrom); err != nil { if err = client.Mail(config.SMTPFrom); err != nil {
return err return err

22
main.go
View File

@@ -27,27 +27,19 @@ func main() {
common.Init() common.Init()
logger.SetupLogger() logger.SetupLogger()
logger.SysLogf("One API %s started", common.Version) logger.SysLogf("One API %s started", common.Version)
if os.Getenv("GIN_MODE") != "debug" {
if os.Getenv("GIN_MODE") != gin.DebugMode {
gin.SetMode(gin.ReleaseMode) gin.SetMode(gin.ReleaseMode)
} }
if config.DebugEnabled { if config.DebugEnabled {
logger.SysLog("running in debug mode") logger.SysLog("running in debug mode")
} }
var err error
// Initialize SQL Database // Initialize SQL Database
model.DB, err = model.InitDB("SQL_DSN") model.InitDB()
if err != nil { model.InitLogDB()
logger.FatalLog("failed to initialize database: " + err.Error())
} var err error
if os.Getenv("LOG_SQL_DSN") != "" {
logger.SysLog("using secondary database for table logs")
model.LOG_DB, err = model.InitDB("LOG_SQL_DSN")
if err != nil {
logger.FatalLog("failed to initialize secondary database: " + err.Error())
}
} else {
model.LOG_DB = model.DB
}
err = model.CreateRootAccountIfNeed() err = model.CreateRootAccountIfNeed()
if err != nil { if err != nil {
logger.FatalLog("database init error: " + err.Error()) logger.FatalLog("database init error: " + err.Error())

View File

@@ -1,6 +1,7 @@
package model package model
import ( import (
"database/sql"
"fmt" "fmt"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
@@ -60,90 +61,156 @@ func CreateRootAccountIfNeed() error {
} }
func chooseDB(envName string) (*gorm.DB, error) { func chooseDB(envName string) (*gorm.DB, error) {
if os.Getenv(envName) != "" { dsn := os.Getenv(envName)
dsn := os.Getenv(envName)
if strings.HasPrefix(dsn, "postgres://") { switch {
// Use PostgreSQL case strings.HasPrefix(dsn, "postgres://"):
logger.SysLog("using PostgreSQL as database") // Use PostgreSQL
common.UsingPostgreSQL = true return openPostgreSQL(dsn)
return gorm.Open(postgres.New(postgres.Config{ case dsn != "":
DSN: dsn,
PreferSimpleProtocol: true, // disables implicit prepared statement usage
}), &gorm.Config{
PrepareStmt: true, // precompile SQL
})
}
// Use MySQL // Use MySQL
logger.SysLog("using MySQL as database") return openMySQL(dsn)
common.UsingMySQL = true default:
return gorm.Open(mysql.Open(dsn), &gorm.Config{ // Use SQLite
PrepareStmt: true, // precompile SQL return openSQLite()
})
} }
// Use SQLite }
logger.SysLog("SQL_DSN not set, using SQLite as database")
common.UsingSQLite = true func openPostgreSQL(dsn string) (*gorm.DB, error) {
config := fmt.Sprintf("?_busy_timeout=%d", common.SQLiteBusyTimeout) logger.SysLog("using PostgreSQL as database")
return gorm.Open(sqlite.Open(common.SQLitePath+config), &gorm.Config{ common.UsingPostgreSQL = true
return gorm.Open(postgres.New(postgres.Config{
DSN: dsn,
PreferSimpleProtocol: true, // disables implicit prepared statement usage
}), &gorm.Config{
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
} }
func InitDB(envName string) (db *gorm.DB, err error) { func openMySQL(dsn string) (*gorm.DB, error) {
db, err = chooseDB(envName) logger.SysLog("using MySQL as database")
if err == nil { common.UsingMySQL = true
if config.DebugSQLEnabled { return gorm.Open(mysql.Open(dsn), &gorm.Config{
db = db.Debug() PrepareStmt: true, // precompile SQL
} })
sqlDB, err := db.DB() }
if err != nil {
return nil, err
}
sqlDB.SetMaxIdleConns(env.Int("SQL_MAX_IDLE_CONNS", 100))
sqlDB.SetMaxOpenConns(env.Int("SQL_MAX_OPEN_CONNS", 1000))
sqlDB.SetConnMaxLifetime(time.Second * time.Duration(env.Int("SQL_MAX_LIFETIME", 60)))
if !config.IsMasterNode { func openSQLite() (*gorm.DB, error) {
return db, err logger.SysLog("SQL_DSN not set, using SQLite as database")
} common.UsingSQLite = true
if common.UsingMySQL { dsn := fmt.Sprintf("%s?_busy_timeout=%d", common.SQLitePath, common.SQLiteBusyTimeout)
_, _ = sqlDB.Exec("DROP INDEX idx_channels_key ON channels;") // TODO: delete this line when most users have upgraded return gorm.Open(sqlite.Open(dsn), &gorm.Config{
} PrepareStmt: true, // precompile SQL
logger.SysLog("database migration started") })
err = db.AutoMigrate(&Channel{}) }
if err != nil {
return nil, err func InitDB() {
} var err error
err = db.AutoMigrate(&Token{}) DB, err = chooseDB("SQL_DSN")
if err != nil { if err != nil {
return nil, err logger.FatalLog("failed to initialize database: " + err.Error())
} return
err = db.AutoMigrate(&User{})
if err != nil {
return nil, err
}
err = db.AutoMigrate(&Option{})
if err != nil {
return nil, err
}
err = db.AutoMigrate(&Redemption{})
if err != nil {
return nil, err
}
err = db.AutoMigrate(&Ability{})
if err != nil {
return nil, err
}
err = db.AutoMigrate(&Log{})
if err != nil {
return nil, err
}
logger.SysLog("database migrated")
return db, err
} else {
logger.FatalLog(err)
} }
return db, err
sqlDB := setDBConns(DB)
if !config.IsMasterNode {
return
}
if common.UsingMySQL {
_, _ = sqlDB.Exec("DROP INDEX idx_channels_key ON channels;") // TODO: delete this line when most users have upgraded
}
logger.SysLog("database migration started")
if err = migrateDB(); err != nil {
logger.FatalLog("failed to migrate database: " + err.Error())
return
}
logger.SysLog("database migrated")
}
func migrateDB() error {
var err error
if err = DB.AutoMigrate(&Channel{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Token{}); err != nil {
return err
}
if err = DB.AutoMigrate(&User{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Option{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Redemption{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Ability{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Log{}); err != nil {
return err
}
if err = DB.AutoMigrate(&Channel{}); err != nil {
return err
}
return nil
}
func InitLogDB() {
if os.Getenv("LOG_SQL_DSN") == "" {
LOG_DB = DB
return
}
logger.SysLog("using secondary database for table logs")
var err error
LOG_DB, err = chooseDB("LOG_SQL_DSN")
if err != nil {
logger.FatalLog("failed to initialize secondary database: " + err.Error())
return
}
setDBConns(LOG_DB)
if !config.IsMasterNode {
return
}
logger.SysLog("secondary database migration started")
err = migrateLOGDB()
if err != nil {
logger.FatalLog("failed to migrate secondary database: " + err.Error())
return
}
logger.SysLog("secondary database migrated")
}
func migrateLOGDB() error {
var err error
if err = LOG_DB.AutoMigrate(&Log{}); err != nil {
return err
}
return nil
}
func setDBConns(db *gorm.DB) *sql.DB {
if config.DebugSQLEnabled {
db = db.Debug()
}
sqlDB, err := db.DB()
if err != nil {
logger.FatalLog("failed to connect database: " + err.Error())
return nil
}
sqlDB.SetMaxIdleConns(env.Int("SQL_MAX_IDLE_CONNS", 100))
sqlDB.SetMaxOpenConns(env.Int("SQL_MAX_OPEN_CONNS", 1000))
sqlDB.SetConnMaxLifetime(time.Second * time.Duration(env.Int("SQL_MAX_LIFETIME", 60)))
return sqlDB
} }
func closeDB(db *gorm.DB) error { func closeDB(db *gorm.DB) error {