Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support ClickHouse log database (#5663)

Browse Source 
* feat: support ClickHouse log database

* feat(log): optimize log deletion process for ClickHouse
This commit is contained in:
Calcium-Ion
2026-06-22 18:41:26 +08:00
committed by GitHub
parent 354d0fedba
commit 6dc4030fdf
25 changed files with 3149 additions and 453 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/main.go
+163 -62
View File
@@ -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