feat: support ClickHouse log database (#5663)
* feat: support ClickHouse log database * feat(log): optimize log deletion process for ClickHouse
This commit is contained in:
+2
-2
@@ -113,7 +113,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if common.UsingSQLite || common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) || common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
err = channelQuery.Order("weight DESC").Find(&abilities).Error
|
||||
} else {
|
||||
err = channelQuery.Order("weight DESC").Find(&abilities).Error
|
||||
@@ -341,7 +341,7 @@ func FixAbility() (int, int, error) {
|
||||
defer fixLock.Unlock()
|
||||
|
||||
// truncate abilities table
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
err := DB.Exec("DELETE FROM abilities").Error
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
|
||||
|
||||
+5
-5
@@ -138,7 +138,7 @@ func NormalizeChannelGroupFilter(group string) string {
|
||||
}
|
||||
|
||||
func channelGroupFilterCondition() string {
|
||||
if common.UsingMySQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
return `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ? ESCAPE '!'`
|
||||
}
|
||||
return `(',' || ` + commonGroupCol + ` || ',') LIKE ? ESCAPE '!'`
|
||||
@@ -381,13 +381,13 @@ func SearchChannels(keyword string, group string, model string, idSort bool, sor
|
||||
modelsCol := "`models`"
|
||||
|
||||
// 如果是 PostgreSQL,使用双引号
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
modelsCol = `"models"`
|
||||
}
|
||||
|
||||
baseURLCol := "`base_url`"
|
||||
// 如果是 PostgreSQL,使用双引号
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
baseURLCol = `"base_url"`
|
||||
}
|
||||
|
||||
@@ -898,13 +898,13 @@ func SearchTags(keyword string, group string, model string, idSort bool) ([]*str
|
||||
modelsCol := "`models`"
|
||||
|
||||
// 如果是 PostgreSQL,使用双引号
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
modelsCol = `"models"`
|
||||
}
|
||||
|
||||
baseURLCol := "`base_url`"
|
||||
// 如果是 PostgreSQL,使用双引号
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
baseURLCol = `"base_url"`
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -82,7 +82,7 @@ func UserCheckin(userId int) (*Checkin, error) {
|
||||
}
|
||||
|
||||
// 根据数据库类型选择不同的策略
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
// SQLite 不支持嵌套事务,使用顺序操作 + 手动回滚
|
||||
return userCheckinWithoutTransaction(checkin, userId, quotaAwarded)
|
||||
}
|
||||
|
||||
+2
-2
@@ -8,9 +8,9 @@ func GetDBTimestamp() int64 {
|
||||
var ts int64
|
||||
var err error
|
||||
switch {
|
||||
case common.UsingPostgreSQL:
|
||||
case common.UsingMainDatabase(common.DatabaseTypePostgreSQL):
|
||||
err = DB.Raw("SELECT EXTRACT(EPOCH FROM NOW())::bigint").Scan(&ts).Error
|
||||
case common.UsingSQLite:
|
||||
case common.UsingMainDatabase(common.DatabaseTypeSQLite):
|
||||
err = DB.Raw("SELECT strftime('%s','now')").Scan(&ts).Error
|
||||
default:
|
||||
err = DB.Raw("SELECT UNIX_TIMESTAMP()").Scan(&ts).Error
|
||||
|
||||
+99
-15
@@ -67,6 +67,27 @@ const (
|
||||
LogTypeLogin = 7
|
||||
)
|
||||
|
||||
func ensureLogRequestId(log *Log) {
|
||||
if log != nil && log.RequestId == "" {
|
||||
log.RequestId = common.NewRequestId()
|
||||
}
|
||||
}
|
||||
|
||||
func createLog(log *Log) error {
|
||||
ensureLogRequestId(log)
|
||||
return LOG_DB.Create(log).Error
|
||||
}
|
||||
|
||||
func clickHouseLogOrder(prefix string) string {
|
||||
return prefix + "created_at desc, " + prefix + "request_id desc"
|
||||
}
|
||||
|
||||
func assignDisplayLogIds(logs []*Log, startIdx int) {
|
||||
for i := range logs {
|
||||
logs[i].Id = startIdx + i + 1
|
||||
}
|
||||
}
|
||||
|
||||
func formatUserLogs(logs []*Log, startIdx int) {
|
||||
for i := range logs {
|
||||
logs[i].ChannelName = ""
|
||||
@@ -81,12 +102,16 @@ func formatUserLogs(logs []*Log, startIdx int) {
|
||||
delete(otherMap, "stream_status")
|
||||
}
|
||||
logs[i].Other = common.MapToJsonStr(otherMap)
|
||||
logs[i].Id = startIdx + i + 1
|
||||
}
|
||||
assignDisplayLogIds(logs, startIdx)
|
||||
}
|
||||
|
||||
func GetLogByTokenId(tokenId int) (logs []*Log, err error) {
|
||||
err = LOG_DB.Model(&Log{}).Where("token_id = ?", tokenId).Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error
|
||||
order := "id desc"
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
order = clickHouseLogOrder("")
|
||||
}
|
||||
err = LOG_DB.Model(&Log{}).Where("token_id = ?", tokenId).Order(order).Limit(common.MaxRecentItems).Find(&logs).Error
|
||||
formatUserLogs(logs, 0)
|
||||
return logs, err
|
||||
}
|
||||
@@ -103,7 +128,7 @@ func RecordLog(userId int, logType int, content string) {
|
||||
Type: logType,
|
||||
Content: content,
|
||||
}
|
||||
err := LOG_DB.Create(log).Error
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
common.SysLog("failed to record log: " + err.Error())
|
||||
}
|
||||
@@ -128,7 +153,7 @@ func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo m
|
||||
}
|
||||
log.Other = common.MapToJsonStr(other)
|
||||
}
|
||||
if err := LOG_DB.Create(log).Error; err != nil {
|
||||
if err := createLog(log); err != nil {
|
||||
common.SysLog("failed to record log: " + err.Error())
|
||||
}
|
||||
}
|
||||
@@ -165,7 +190,7 @@ func RecordLoginLog(userId int, username string, content string, ip string, acti
|
||||
Ip: ip,
|
||||
Other: common.MapToJsonStr(other),
|
||||
}
|
||||
if err := LOG_DB.Create(log).Error; err != nil {
|
||||
if err := createLog(log); err != nil {
|
||||
common.SysLog("failed to record login log: " + err.Error())
|
||||
}
|
||||
}
|
||||
@@ -196,7 +221,7 @@ func RecordOperationAuditLog(logUserId int, content string, ip string, action st
|
||||
Ip: ip,
|
||||
Other: common.MapToJsonStr(other),
|
||||
}
|
||||
if err := LOG_DB.Create(log).Error; err != nil {
|
||||
if err := createLog(log); err != nil {
|
||||
common.SysLog("failed to record operation audit log: " + err.Error())
|
||||
}
|
||||
}
|
||||
@@ -223,7 +248,7 @@ func RecordTopupLog(userId int, content string, callerIp string, paymentMethod s
|
||||
Ip: callerIp,
|
||||
Other: common.MapToJsonStr(other),
|
||||
}
|
||||
err := LOG_DB.Create(log).Error
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
common.SysLog("failed to record topup log: " + err.Error())
|
||||
}
|
||||
@@ -269,7 +294,7 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string,
|
||||
UpstreamRequestId: upstreamRequestId,
|
||||
Other: otherStr,
|
||||
}
|
||||
err := LOG_DB.Create(log).Error
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
logger.LogError(c, "failed to record log: "+err.Error())
|
||||
}
|
||||
@@ -333,7 +358,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
|
||||
UpstreamRequestId: upstreamRequestId,
|
||||
Other: otherStr,
|
||||
}
|
||||
err := LOG_DB.Create(log).Error
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
logger.LogError(c, "failed to record log: "+err.Error())
|
||||
}
|
||||
@@ -393,7 +418,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
|
||||
Group: params.Group,
|
||||
Other: common.MapToJsonStr(params.Other),
|
||||
}
|
||||
err := LOG_DB.Create(log).Error
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
common.SysLog("failed to record task billing log: " + err.Error())
|
||||
}
|
||||
@@ -453,10 +478,17 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
err = tx.Order("logs.created_at desc, logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error
|
||||
order := "logs.created_at desc, logs.id desc"
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
order = clickHouseLogOrder("logs.")
|
||||
}
|
||||
err = tx.Order(order).Limit(num).Offset(startIdx).Find(&logs).Error
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
assignDisplayLogIds(logs, startIdx)
|
||||
}
|
||||
|
||||
channelIds := types.NewSet[int]()
|
||||
for _, log := range logs {
|
||||
@@ -537,7 +569,11 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int
|
||||
common.SysError("failed to count user logs: " + err.Error())
|
||||
return nil, 0, errors.New("查询日志失败")
|
||||
}
|
||||
err = tx.Order("logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error
|
||||
order := "logs.id desc"
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
order = clickHouseLogOrder("logs.")
|
||||
}
|
||||
err = tx.Order(order).Limit(num).Offset(startIdx).Find(&logs).Error
|
||||
if err != nil {
|
||||
common.SysError("failed to search user logs: " + err.Error())
|
||||
return nil, 0, errors.New("查询日志失败")
|
||||
@@ -554,10 +590,10 @@ type Stat struct {
|
||||
}
|
||||
|
||||
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
|
||||
tx := LOG_DB.Table("logs").Select("sum(quota) quota")
|
||||
tx := LOG_DB.Table("logs").Select("COALESCE(sum(quota), 0) quota")
|
||||
|
||||
// 为rpm和tpm创建单独的查询
|
||||
rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm")
|
||||
rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, COALESCE(sum(prompt_tokens), 0) + COALESCE(sum(completion_tokens), 0) tpm")
|
||||
|
||||
if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil {
|
||||
return stat, err
|
||||
@@ -610,7 +646,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) {
|
||||
tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
|
||||
tx := LOG_DB.Table("logs").Select("COALESCE(sum(prompt_tokens), 0) + COALESCE(sum(completion_tokens), 0)")
|
||||
if username != "" {
|
||||
tx = tx.Where("username = ?", username)
|
||||
}
|
||||
@@ -631,6 +667,54 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa
|
||||
}
|
||||
|
||||
func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
|
||||
var total int64 = 0
|
||||
|
||||
for {
|
||||
if nil != ctx.Err() {
|
||||
return total, ctx.Err()
|
||||
}
|
||||
|
||||
var batchCount int64
|
||||
if err := LOG_DB.WithContext(ctx).Raw(`
|
||||
SELECT count() FROM (
|
||||
SELECT created_at, request_id
|
||||
FROM logs
|
||||
WHERE created_at < ?
|
||||
ORDER BY created_at ASC, request_id ASC
|
||||
LIMIT ?
|
||||
)`, targetTimestamp, limit).Scan(&batchCount).Error; err != nil {
|
||||
return total, err
|
||||
}
|
||||
if batchCount == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
if err := LOG_DB.WithContext(ctx).Exec(`
|
||||
ALTER TABLE logs DELETE WHERE (created_at, request_id) IN (
|
||||
SELECT created_at, request_id
|
||||
FROM logs
|
||||
WHERE created_at < ?
|
||||
ORDER BY created_at ASC, request_id ASC
|
||||
LIMIT ?
|
||||
) SETTINGS mutations_sync = 1`, targetTimestamp, limit).Error; err != nil {
|
||||
return total, err
|
||||
}
|
||||
|
||||
total += batchCount
|
||||
|
||||
if batchCount < int64(limit) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return total, nil
|
||||
}
|
||||
|
||||
var total int64 = 0
|
||||
|
||||
for {
|
||||
|
||||
+163
-62
@@ -3,6 +3,7 @@ package model
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -12,6 +13,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/driver/clickhouse"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
@@ -27,7 +29,7 @@ var logGroupCol string
|
||||
|
||||
func initCol() {
|
||||
// init common column names
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
commonGroupCol = `"group"`
|
||||
commonKeyCol = `"key"`
|
||||
commonTrueVal = "true"
|
||||
@@ -38,27 +40,14 @@ func initCol() {
|
||||
commonTrueVal = "1"
|
||||
commonFalseVal = "0"
|
||||
}
|
||||
if os.Getenv("LOG_SQL_DSN") != "" {
|
||||
switch common.LogSqlType {
|
||||
case common.DatabaseTypePostgreSQL:
|
||||
logGroupCol = `"group"`
|
||||
logKeyCol = `"key"`
|
||||
default:
|
||||
logGroupCol = commonGroupCol
|
||||
logKeyCol = commonKeyCol
|
||||
}
|
||||
} else {
|
||||
// LOG_SQL_DSN 为空时,日志数据库与主数据库相同
|
||||
if common.UsingPostgreSQL {
|
||||
logGroupCol = `"group"`
|
||||
logKeyCol = `"key"`
|
||||
} else {
|
||||
logGroupCol = commonGroupCol
|
||||
logKeyCol = commonKeyCol
|
||||
}
|
||||
switch common.LogDatabaseType() {
|
||||
case common.DatabaseTypePostgreSQL:
|
||||
logGroupCol = `"group"`
|
||||
logKeyCol = `"key"`
|
||||
default:
|
||||
logGroupCol = "`group`"
|
||||
logKeyCol = "`key`"
|
||||
}
|
||||
// log sql type and database type
|
||||
//common.SysLog("Using Log SQL Type: " + common.LogSqlType)
|
||||
}
|
||||
|
||||
var DB *gorm.DB
|
||||
@@ -115,37 +104,56 @@ func CheckSetup() {
|
||||
}
|
||||
}
|
||||
|
||||
func chooseDB(envName string, isLog bool) (*gorm.DB, error) {
|
||||
defer func() {
|
||||
initCol()
|
||||
}()
|
||||
func isClickHouseDSN(dsn string) bool {
|
||||
return strings.HasPrefix(dsn, "clickhouse://") ||
|
||||
strings.HasPrefix(dsn, "tcp://") ||
|
||||
strings.HasPrefix(dsn, "http://") ||
|
||||
strings.HasPrefix(dsn, "https://")
|
||||
}
|
||||
|
||||
func normalizeClickHouseDSN(dsn string) string {
|
||||
parsed, err := url.Parse(dsn)
|
||||
if err != nil || parsed.Scheme != "https" {
|
||||
return dsn
|
||||
}
|
||||
query := parsed.Query()
|
||||
if _, ok := query["secure"]; !ok {
|
||||
query.Set("secure", "true")
|
||||
parsed.RawQuery = query.Encode()
|
||||
}
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
func chooseDB(envName string, isLog bool) (*gorm.DB, common.DatabaseType, error) {
|
||||
dsn := os.Getenv(envName)
|
||||
if dsn != "" {
|
||||
if isClickHouseDSN(dsn) {
|
||||
if !isLog {
|
||||
return nil, "", fmt.Errorf("%s does not support ClickHouse; use SQLite, MySQL, or PostgreSQL for the primary database and LOG_SQL_DSN for ClickHouse logs", envName)
|
||||
}
|
||||
common.SysLog("using ClickHouse as log database")
|
||||
db, err := gorm.Open(clickhouse.Open(normalizeClickHouseDSN(dsn)), &gorm.Config{
|
||||
PrepareStmt: false,
|
||||
})
|
||||
return db, common.DatabaseTypeClickHouse, err
|
||||
}
|
||||
if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
|
||||
// Use PostgreSQL
|
||||
common.SysLog("using PostgreSQL as database")
|
||||
if !isLog {
|
||||
common.UsingPostgreSQL = true
|
||||
} else {
|
||||
common.LogSqlType = common.DatabaseTypePostgreSQL
|
||||
}
|
||||
return gorm.Open(postgres.New(postgres.Config{
|
||||
db, err := gorm.Open(postgres.New(postgres.Config{
|
||||
DSN: dsn,
|
||||
PreferSimpleProtocol: true, // disables implicit prepared statement usage
|
||||
}), &gorm.Config{
|
||||
PrepareStmt: true, // precompile SQL
|
||||
})
|
||||
return db, common.DatabaseTypePostgreSQL, err
|
||||
}
|
||||
if strings.HasPrefix(dsn, "local") {
|
||||
common.SysLog("SQL_DSN not set, using SQLite as database")
|
||||
if !isLog {
|
||||
common.UsingSQLite = true
|
||||
} else {
|
||||
common.LogSqlType = common.DatabaseTypeSQLite
|
||||
}
|
||||
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
|
||||
db, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
|
||||
PrepareStmt: true, // precompile SQL
|
||||
})
|
||||
return db, common.DatabaseTypeSQLite, err
|
||||
}
|
||||
// Use MySQL
|
||||
common.SysLog("using MySQL as database")
|
||||
@@ -157,32 +165,33 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) {
|
||||
dsn += "?parseTime=true"
|
||||
}
|
||||
}
|
||||
if !isLog {
|
||||
common.UsingMySQL = true
|
||||
} else {
|
||||
common.LogSqlType = common.DatabaseTypeMySQL
|
||||
}
|
||||
return gorm.Open(mysql.Open(dsn), &gorm.Config{
|
||||
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
|
||||
PrepareStmt: true, // precompile SQL
|
||||
})
|
||||
return db, common.DatabaseTypeMySQL, err
|
||||
}
|
||||
// Use SQLite
|
||||
common.SysLog("SQL_DSN not set, using SQLite as database")
|
||||
common.UsingSQLite = true
|
||||
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
|
||||
db, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
|
||||
PrepareStmt: true, // precompile SQL
|
||||
})
|
||||
return db, common.DatabaseTypeSQLite, err
|
||||
}
|
||||
|
||||
func InitDB() (err error) {
|
||||
db, err := chooseDB("SQL_DSN", false)
|
||||
db, dbType, err := chooseDB("SQL_DSN", false)
|
||||
if err == nil {
|
||||
common.SetMainDatabaseType(dbType)
|
||||
if os.Getenv("LOG_SQL_DSN") == "" {
|
||||
common.SetLogDatabaseType(dbType)
|
||||
}
|
||||
initCol()
|
||||
if common.DebugEnabled {
|
||||
db = db.Debug()
|
||||
}
|
||||
DB = db
|
||||
// MySQL charset/collation startup check: ensure Chinese-capable charset
|
||||
if common.UsingMySQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
if err := checkMySQLChineseSupport(DB); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
@@ -198,7 +207,7 @@ func InitDB() (err error) {
|
||||
if !common.IsMasterNode {
|
||||
return nil
|
||||
}
|
||||
if common.UsingMySQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
//_, _ = sqlDB.Exec("ALTER TABLE channels MODIFY model_mapping TEXT;") // TODO: delete this line when most users have upgraded
|
||||
}
|
||||
common.SysLog("database migration started")
|
||||
@@ -213,16 +222,20 @@ func InitDB() (err error) {
|
||||
func InitLogDB() (err error) {
|
||||
if os.Getenv("LOG_SQL_DSN") == "" {
|
||||
LOG_DB = DB
|
||||
common.SetLogDatabaseType(common.MainDatabaseType())
|
||||
initCol()
|
||||
return
|
||||
}
|
||||
db, err := chooseDB("LOG_SQL_DSN", true)
|
||||
db, dbType, err := chooseDB("LOG_SQL_DSN", true)
|
||||
if err == nil {
|
||||
common.SetLogDatabaseType(dbType)
|
||||
initCol()
|
||||
if common.DebugEnabled {
|
||||
db = db.Debug()
|
||||
}
|
||||
LOG_DB = db
|
||||
// If log DB is MySQL, also ensure Chinese-capable charset
|
||||
if common.LogSqlType == common.DatabaseTypeMySQL {
|
||||
if common.UsingLogDatabase(common.DatabaseTypeMySQL) {
|
||||
if err := checkMySQLChineseSupport(LOG_DB); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
@@ -285,7 +298,7 @@ func migrateDB() error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
if err := ensureSubscriptionPlanTableSQLite(); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -354,7 +367,7 @@ func migrateDBFast() error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
if err := ensureSubscriptionPlanTableSQLite(); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -368,11 +381,99 @@ func migrateDBFast() error {
|
||||
}
|
||||
|
||||
func migrateLOGDB() error {
|
||||
var err error
|
||||
if err = LOG_DB.AutoMigrate(&Log{}); err != nil {
|
||||
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
|
||||
return migrateClickHouseLogDB()
|
||||
}
|
||||
return LOG_DB.AutoMigrate(&Log{})
|
||||
}
|
||||
|
||||
func migrateClickHouseLogDB() error {
|
||||
ttlDays := clickHouseLogTTLDays()
|
||||
if err := LOG_DB.Exec(clickHouseLogCreateTableSQL(ttlDays)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
return syncClickHouseLogTTL(ttlDays)
|
||||
}
|
||||
|
||||
func clickHouseLogTTLDays() int {
|
||||
ttlDays := common.GetEnvOrDefault("LOG_SQL_CLICKHOUSE_TTL_DAYS", 0)
|
||||
if ttlDays < 0 {
|
||||
return 0
|
||||
}
|
||||
return ttlDays
|
||||
}
|
||||
|
||||
func clickHouseLogTTLExpression(ttlDays int) string {
|
||||
if ttlDays <= 0 {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprintf("toDateTime(created_at) + INTERVAL %d DAY DELETE", ttlDays)
|
||||
}
|
||||
|
||||
func clickHouseLogTTLClause(ttlDays int) string {
|
||||
expression := clickHouseLogTTLExpression(ttlDays)
|
||||
if expression == "" {
|
||||
return ""
|
||||
}
|
||||
return "\nTTL " + expression
|
||||
}
|
||||
|
||||
func clickHouseLogCreateTableSQL(ttlDays int) string {
|
||||
return fmt.Sprintf(`
|
||||
CREATE TABLE IF NOT EXISTS logs (
|
||||
id Int64 DEFAULT 0,
|
||||
user_id Int32 DEFAULT 0,
|
||||
created_at Int64 DEFAULT 0,
|
||||
type Int32 DEFAULT 0,
|
||||
content String DEFAULT '',
|
||||
username String DEFAULT '',
|
||||
token_name String DEFAULT '',
|
||||
model_name String DEFAULT '',
|
||||
quota Int32 DEFAULT 0,
|
||||
prompt_tokens Int32 DEFAULT 0,
|
||||
completion_tokens Int32 DEFAULT 0,
|
||||
use_time Int32 DEFAULT 0,
|
||||
is_stream UInt8 DEFAULT 0,
|
||||
channel_id Int32 DEFAULT 0,
|
||||
token_id Int32 DEFAULT 0,
|
||||
`+"`group`"+` String DEFAULT '',
|
||||
ip String DEFAULT '',
|
||||
request_id String DEFAULT '',
|
||||
upstream_request_id String DEFAULT '',
|
||||
other String DEFAULT ''
|
||||
)
|
||||
ENGINE = MergeTree()
|
||||
PARTITION BY toYYYYMM(toDateTime(created_at))
|
||||
ORDER BY (created_at, request_id)%s`, clickHouseLogTTLClause(ttlDays))
|
||||
}
|
||||
|
||||
func syncClickHouseLogTTL(ttlDays int) error {
|
||||
expression := clickHouseLogTTLExpression(ttlDays)
|
||||
if expression != "" {
|
||||
return LOG_DB.Exec("ALTER TABLE logs MODIFY TTL " + expression).Error
|
||||
}
|
||||
|
||||
hasTTL, err := clickHouseLogTableHasTTL()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !hasTTL {
|
||||
return nil
|
||||
}
|
||||
return LOG_DB.Exec("ALTER TABLE logs REMOVE TTL").Error
|
||||
}
|
||||
|
||||
func clickHouseLogTableHasTTL() (bool, error) {
|
||||
var createTableSQL string
|
||||
if err := LOG_DB.Raw("SHOW CREATE TABLE logs").Scan(&createTableSQL).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return clickHouseCreateTableHasTTL(createTableSQL), nil
|
||||
}
|
||||
|
||||
func clickHouseCreateTableHasTTL(createTableSQL string) bool {
|
||||
upperSQL := strings.ToUpper(createTableSQL)
|
||||
return strings.Contains(upperSQL, "\nTTL ") || strings.Contains(upperSQL, " TTL ")
|
||||
}
|
||||
|
||||
type sqliteColumnDef struct {
|
||||
@@ -381,7 +482,7 @@ type sqliteColumnDef struct {
|
||||
}
|
||||
|
||||
func ensureSubscriptionPlanTableSQLite() error {
|
||||
if !common.UsingSQLite {
|
||||
if !common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
return nil
|
||||
}
|
||||
tableName := "subscription_plans"
|
||||
@@ -463,7 +564,7 @@ PRIMARY KEY (` + "`id`" + `)
|
||||
// This is safe to run multiple times - it checks the column type first
|
||||
func migrateTokenModelLimitsToText() error {
|
||||
// SQLite uses type affinity, so TEXT and VARCHAR are effectively the same — no migration needed
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -479,7 +580,7 @@ func migrateTokenModelLimitsToText() error {
|
||||
}
|
||||
|
||||
var alterSQL string
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
var dataType string
|
||||
if err := DB.Raw(`SELECT data_type FROM information_schema.columns
|
||||
WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`,
|
||||
@@ -489,7 +590,7 @@ func migrateTokenModelLimitsToText() error {
|
||||
return nil
|
||||
}
|
||||
alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE text`, tableName, columnName)
|
||||
} else if common.UsingMySQL {
|
||||
} else if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
var columnType string
|
||||
if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
|
||||
WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
|
||||
@@ -517,7 +618,7 @@ func migrateTokenModelLimitsToText() error {
|
||||
func migrateSubscriptionPlanPriceAmount() {
|
||||
// SQLite doesn't support ALTER COLUMN, and its type affinity handles this automatically
|
||||
// Skip early to avoid GORM parsing the existing table DDL which may cause issues
|
||||
if common.UsingSQLite {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -535,7 +636,7 @@ func migrateSubscriptionPlanPriceAmount() {
|
||||
}
|
||||
|
||||
var alterSQL string
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
// PostgreSQL: Check if already decimal/numeric
|
||||
var dataType string
|
||||
if err := DB.Raw(`SELECT data_type FROM information_schema.columns
|
||||
@@ -547,7 +648,7 @@ func migrateSubscriptionPlanPriceAmount() {
|
||||
}
|
||||
alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE decimal(10,6) USING %s::decimal(10,6)`,
|
||||
tableName, columnName, columnName)
|
||||
} else if common.UsingMySQL {
|
||||
} else if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
// MySQL: Check if already decimal
|
||||
var columnType string
|
||||
if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
|
||||
|
||||
+1
-1
@@ -122,7 +122,7 @@ func Redeem(key string, userId int) (quota int, err error) {
|
||||
redemption := &Redemption{}
|
||||
|
||||
keyCol := "`key`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
keyCol = `"key"`
|
||||
}
|
||||
common.RandomSleep()
|
||||
|
||||
@@ -555,7 +555,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP
|
||||
return errors.New("tradeNo is empty")
|
||||
}
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
var logUserId int
|
||||
@@ -663,7 +663,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err
|
||||
return errors.New("tradeNo is empty")
|
||||
}
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
|
||||
@@ -22,7 +22,7 @@ func TestMain(m *testing.M) {
|
||||
DB = db
|
||||
LOG_DB = db
|
||||
|
||||
common.UsingSQLite = true
|
||||
common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
|
||||
common.RedisEnabled = false
|
||||
common.BatchUpdateEnabled = false
|
||||
common.LogConsumeEnabled = true
|
||||
|
||||
+6
-6
@@ -85,7 +85,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta
|
||||
}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
@@ -115,7 +115,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error
|
||||
topUp := &TopUp{}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
@@ -323,7 +323,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error {
|
||||
}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
@@ -398,7 +398,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
|
||||
topUp := &TopUp{}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
@@ -473,7 +473,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) {
|
||||
topUp := &TopUp{}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
@@ -536,7 +536,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) {
|
||||
topUp := &TopUp{}
|
||||
|
||||
refCol := "`trade_no`"
|
||||
if common.UsingPostgreSQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
|
||||
refCol = `"trade_no"`
|
||||
}
|
||||
|
||||
|
||||
@@ -49,7 +49,7 @@ func GetRankingQuotaBuckets(startTime int64, endTime int64, bucketSize int64) ([
|
||||
}
|
||||
|
||||
func rankingBucketExpr(bucketSize int64) string {
|
||||
if common.UsingMySQL {
|
||||
if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
|
||||
return fmt.Sprintf("FLOOR(created_at / %d) * %d", bucketSize, bucketSize)
|
||||
}
|
||||
return fmt.Sprintf("(created_at / %d) * %d", bucketSize, bucketSize)
|
||||
|
||||
Reference in New Issue
Block a user