Skip to content

phsl / new-api

Go to a project
Commit 6383d13c by Calcium-Ion Committed by GitHub
Options

Merge pull request #432 from xixingya/feat-add-logdb 

Feature: Support Log DB
parents a7c14a26 3a6a30a8
Showing with 134 additions and 60 deletions
main.go
...@@ -42,6 +42,11 @@ func main() { ...@@ -42,6 +42,11 @@ func main() {
if err != nil { if err != nil {
common.FatalLog("failed to initialize database: " + err.Error()) common.FatalLog("failed to initialize database: " + err.Error())
} }
// Initialize SQL Database
err = model.InitLogDB()
if err != nil {
common.FatalLog("failed to initialize database: " + err.Error())
}
defer func() { defer func() {
err := model.CloseDB() err := model.CloseDB()
if err != nil { if err != nil {
......
model/log.go
...@@ -38,7 +38,7 @@ const ( ...@@ -38,7 +38,7 @@ const (
) )
func GetLogByKey(key string) (logs []*Log, err error) { func GetLogByKey(key string) (logs []*Log, err error) {
err = DB.Joins("left join tokens on tokens.id = logs.token_id").Where("tokens.key = ?", strings.TrimPrefix(key, "sk-")).Find(&logs).Error err = LOG_DB.Joins("left join tokens on tokens.id = logs.token_id").Where("tokens.key = ?", strings.TrimPrefix(key, "sk-")).Find(&logs).Error
return logs, err return logs, err
} }
...@@ -54,7 +54,7 @@ func RecordLog(userId int, logType int, content string) { ...@@ -54,7 +54,7 @@ func RecordLog(userId int, logType int, content string) {
Type: logType, Type: logType,
Content: content, Content: content,
} }
err := DB.Create(log).Error err := LOG_DB.Create(log).Error
if err != nil { if err != nil {
common.SysError("failed to record log: " + err.Error()) common.SysError("failed to record log: " + err.Error())
} }
...@@ -84,7 +84,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke ...@@ -84,7 +84,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
IsStream: isStream, IsStream: isStream,
Other: otherStr, Other: otherStr,
} }
err := DB.Create(log).Error err := LOG_DB.Create(log).Error
if err != nil { if err != nil {
common.LogError(ctx, "failed to record log: "+err.Error()) common.LogError(ctx, "failed to record log: "+err.Error())
} }
...@@ -98,9 +98,9 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke ...@@ -98,9 +98,9 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) { func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) {
var tx *gorm.DB var tx *gorm.DB
if logType == LogTypeUnknown { if logType == LogTypeUnknown {
tx = DB tx = LOG_DB
} else { } else {
tx = DB.Where("type = ?", logType) tx = LOG_DB.Where("type = ?", logType)
} }
if modelName != "" { if modelName != "" {
tx = tx.Where("model_name like ?", modelName) tx = tx.Where("model_name like ?", modelName)
...@@ -127,9 +127,9 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName ...@@ -127,9 +127,9 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int) (logs []*Log, err error) { func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int) (logs []*Log, err error) {
var tx *gorm.DB var tx *gorm.DB
if logType == LogTypeUnknown { if logType == LogTypeUnknown {
tx = DB.Where("user_id = ?", userId) tx = LOG_DB.Where("user_id = ?", userId)
} else { } else {
tx = DB.Where("user_id = ? and type = ?", userId, logType) tx = LOG_DB.Where("user_id = ? and type = ?", userId, logType)
} }
if modelName != "" { if modelName != "" {
tx = tx.Where("model_name like ?", modelName) tx = tx.Where("model_name like ?", modelName)
...@@ -157,12 +157,12 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int ...@@ -157,12 +157,12 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int
} }
func SearchAllLogs(keyword string) (logs []*Log, err error) { func SearchAllLogs(keyword string) (logs []*Log, err error) {
err = DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error err = LOG_DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error
return logs, err return logs, err
} }
func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) { func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) {
err = DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(common.MaxRecentItems).Omit("id").Find(&logs).Error err = LOG_DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(common.MaxRecentItems).Omit("id").Find(&logs).Error
return logs, err return logs, err
} }
...@@ -173,10 +173,10 @@ type Stat struct { ...@@ -173,10 +173,10 @@ type Stat struct {
} }
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (stat Stat) { func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (stat Stat) {
tx := DB.Table("logs").Select("sum(quota) quota") tx := LOG_DB.Table("logs").Select("sum(quota) quota")
// 为rpm和tpm创建单独的查询 // 为rpm和tpm创建单独的查询
rpmTpmQuery := DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm") rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm")
if username != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)
...@@ -215,7 +215,7 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa ...@@ -215,7 +215,7 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa
} }
func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) { func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) {
tx := DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)") tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
if username != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)
} }
...@@ -236,6 +236,6 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa ...@@ -236,6 +236,6 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa
} }
func DeleteOldLog(targetTimestamp int64) (int64, error) { func DeleteOldLog(targetTimestamp int64) (int64, error) {
result := DB.Where("created_at < ?", targetTimestamp).Delete(&Log{}) result := LOG_DB.Where("created_at < ?", targetTimestamp).Delete(&Log{})
return result.RowsAffected, result.Error return result.RowsAffected, result.Error
} }
model/main.go
...@@ -15,6 +15,8 @@ import ( ...@@ -15,6 +15,8 @@ import (
var DB *gorm.DB var DB *gorm.DB
var LOG_DB *gorm.DB
func createRootAccountIfNeed() error { func createRootAccountIfNeed() error {
var user User var user User
//if user.Status != common.UserStatusEnabled { //if user.Status != common.UserStatusEnabled {
...@@ -38,9 +40,9 @@ func createRootAccountIfNeed() error { ...@@ -38,9 +40,9 @@ func createRootAccountIfNeed() error {
return nil return nil
} }
func chooseDB() (*gorm.DB, error) { func chooseDB(envName string) (*gorm.DB, error) {
if os.Getenv("SQL_DSN") != "" { dsn := os.Getenv(envName)
dsn := os.Getenv("SQL_DSN") if dsn != "" {
if strings.HasPrefix(dsn, "postgres://") { if strings.HasPrefix(dsn, "postgres://") {
// Use PostgreSQL // Use PostgreSQL
common.SysLog("using PostgreSQL as database") common.SysLog("using PostgreSQL as database")
...@@ -52,6 +54,13 @@ func chooseDB() (*gorm.DB, error) { ...@@ -52,6 +54,13 @@ func chooseDB() (*gorm.DB, error) {
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
} }
if strings.HasPrefix(dsn, "local") {
common.SysLog("SQL_DSN not set, using SQLite as database")
common.UsingSQLite = true
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
PrepareStmt: true, // precompile SQL
})
}
// Use MySQL // Use MySQL
common.SysLog("using MySQL as database") common.SysLog("using MySQL as database")
// check parseTime // check parseTime
...@@ -76,7 +85,7 @@ func chooseDB() (*gorm.DB, error) { ...@@ -76,7 +85,7 @@ func chooseDB() (*gorm.DB, error) {
} }
func InitDB() (err error) { func InitDB() (err error) {
db, err := chooseDB() db, err := chooseDB("SQL_DSN")
if err == nil { if err == nil {
if common.DebugEnabled { if common.DebugEnabled {
db = db.Debug() db = db.Debug()
...@@ -100,52 +109,44 @@ func InitDB() (err error) { ...@@ -100,52 +109,44 @@ func InitDB() (err error) {
// _, _ = sqlDB.Exec("ALTER TABLE midjourneys MODIFY status VARCHAR(20);") // TODO: delete this line when most users have upgraded // _, _ = sqlDB.Exec("ALTER TABLE midjourneys MODIFY status VARCHAR(20);") // TODO: delete this line when most users have upgraded
//} //}
common.SysLog("database migration started") common.SysLog("database migration started")
err = db.AutoMigrate(&Channel{}) err = migrateDB()
if err != nil { return err
return err } else {
} common.FatalLog(err)
err = db.AutoMigrate(&Token{}) }
if err != nil { return err
return err }
}
err = db.AutoMigrate(&User{}) func InitLogDB() (err error) {
if err != nil { if os.Getenv("LOG_SQL_DSN") == "" {
return err LOG_DB = DB
} return
err = db.AutoMigrate(&Option{}) }
if err != nil { db, err := chooseDB("LOG_SQL_DSN")
return err if err == nil {
} if common.DebugEnabled {
err = db.AutoMigrate(&Redemption{}) db = db.Debug()
if err != nil {
return err
}
err = db.AutoMigrate(&Ability{})
if err != nil {
return err
}
err = db.AutoMigrate(&Log{})
if err != nil {
return err
}
err = db.AutoMigrate(&Midjourney{})
if err != nil {
return err
}
err = db.AutoMigrate(&TopUp{})
if err != nil {
return err
} }
err = db.AutoMigrate(&QuotaData{}) LOG_DB = db
sqlDB, err := LOG_DB.DB()
if err != nil { if err != nil {
return err return err
} }
err = db.AutoMigrate(&Task{}) sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100))
if err != nil { sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000))
return err sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60)))
if !common.IsMasterNode {
return nil
} }
common.SysLog("database migrated") //if common.UsingMySQL {
err = createRootAccountIfNeed() // _, _ = sqlDB.Exec("DROP INDEX idx_channels_key ON channels;") // TODO: delete this line when most users have upgraded
// _, _ = sqlDB.Exec("ALTER TABLE midjourneys MODIFY action VARCHAR(40);") // TODO: delete this line when most users have upgraded
// _, _ = sqlDB.Exec("ALTER TABLE midjourneys MODIFY progress VARCHAR(30);") // TODO: delete this line when most users have upgraded
// _, _ = sqlDB.Exec("ALTER TABLE midjourneys MODIFY status VARCHAR(20);") // TODO: delete this line when most users have upgraded
//}
common.SysLog("database migration started")
err = migrateLOGDB()
return err return err
} else { } else {
common.FatalLog(err) common.FatalLog(err)
...@@ -153,8 +154,66 @@ func InitDB() (err error) { ...@@ -153,8 +154,66 @@ func InitDB() (err error) {
return err return err
} }
func CloseDB() error { func migrateDB() error {
sqlDB, err := DB.DB() err := DB.AutoMigrate(&Channel{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Token{})
if err != nil {
return err
}
err = DB.AutoMigrate(&User{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Option{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Redemption{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Ability{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Log{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Midjourney{})
if err != nil {
return err
}
err = DB.AutoMigrate(&TopUp{})
if err != nil {
return err
}
err = DB.AutoMigrate(&QuotaData{})
if err != nil {
return err
}
err = DB.AutoMigrate(&Task{})
if err != nil {
return err
}
common.SysLog("database migrated")
err = createRootAccountIfNeed()
return err
}
func migrateLOGDB() error {
var err error
if err = LOG_DB.AutoMigrate(&Log{}); err != nil {
return err
}
return nil
}
func closeDB(db *gorm.DB) error {
sqlDB, err := db.DB()
if err != nil { if err != nil {
return err return err
} }
...@@ -162,6 +221,16 @@ func CloseDB() error { ...@@ -162,6 +221,16 @@ func CloseDB() error {
return err return err
} }
func CloseDB() error {
if LOG_DB != DB {
err := closeDB(LOG_DB)
if err != nil {
return err
}
}
return closeDB(DB)
}
var ( var (
lastPingTime time.Time lastPingTime time.Time
pingMutex sync.Mutex pingMutex sync.Mutex
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment