Skip to content

phsl / new-api

Go to a project
Commit 6dc4030f by Calcium-Ion Committed by GitHub
Options

feat: support ClickHouse log database (#5663) 

* feat: support ClickHouse log database

* feat(log): optimize log deletion process for ClickHouse
parent 354d0fed
Showing with 406 additions and 445 deletions
AGENTS.md
...@@ -82,7 +82,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag ...@@ -82,7 +82,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag
- PostgreSQL uses `"column"` quoting, while MySQL/SQLite use `` `column` ``. - PostgreSQL uses `"column"` quoting, while MySQL/SQLite use `` `column` ``.
- Use `commonGroupCol`, `commonKeyCol` from `model/main.go` for reserved-word columns like `group` and `key`. - Use `commonGroupCol`, `commonKeyCol` from `model/main.go` for reserved-word columns like `group` and `key`.
- Use `commonTrueVal`/`commonFalseVal` for boolean values. - Use `commonTrueVal`/`commonFalseVal` for boolean values.
- Use `common.UsingPostgreSQL`, `common.UsingSQLite`, and `common.UsingMySQL` flags for DB-specific branches. - Use `common.UsingMainDatabase(...)` for primary database branches and `common.UsingLogDatabase(...)` for log database branches.
- Do not use database-specific features without cross-DB fallback, including MySQL-only functions, PostgreSQL-only operators, SQLite-unsupported `ALTER COLUMN`, or database-specific JSON column types without a `TEXT` fallback. - Do not use database-specific features without cross-DB fallback, including MySQL-only functions, PostgreSQL-only operators, SQLite-unsupported `ALTER COLUMN`, or database-specific JSON column types without a `TEXT` fallback.
- Migrations must work on all three databases. For SQLite, use `ALTER TABLE ... ADD COLUMN` instead of `ALTER COLUMN` (see `model/main.go` for patterns). - Migrations must work on all three databases. For SQLite, use `ALTER TABLE ... ADD COLUMN` instead of `ALTER COLUMN` (see `model/main.go` for patterns).
- Avoid GORM boolean default tags such as `gorm:"default:true"` when the default is a business rule already enforced by code. MySQL and PostgreSQL can normalize boolean defaults differently, causing GORM `AutoMigrate` to repeatedly issue `ALTER TABLE` on restart. Prefer setting these defaults in request/model normalization, hooks, constructors, or service logic; do not replace `default:true` with `default:1` unless the behavior is verified across SQLite, MySQL, and PostgreSQL. - Avoid GORM boolean default tags such as `gorm:"default:true"` when the default is a business rule already enforced by code. MySQL and PostgreSQL can normalize boolean defaults differently, causing GORM `AutoMigrate` to repeatedly issue `ALTER TABLE` on restart. Prefer setting these defaults in request/model normalization, hooks, constructors, or service logic; do not replace `default:true` with `default:1` unless the behavior is verified across SQLite, MySQL, and PostgreSQL.
......
CLAUDE.md
...@@ -82,7 +82,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag ...@@ -82,7 +82,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag
- PostgreSQL uses `"column"` quoting, while MySQL/SQLite use `` `column` ``. - PostgreSQL uses `"column"` quoting, while MySQL/SQLite use `` `column` ``.
- Use `commonGroupCol`, `commonKeyCol` from `model/main.go` for reserved-word columns like `group` and `key`. - Use `commonGroupCol`, `commonKeyCol` from `model/main.go` for reserved-word columns like `group` and `key`.
- Use `commonTrueVal`/`commonFalseVal` for boolean values. - Use `commonTrueVal`/`commonFalseVal` for boolean values.
- Use `common.UsingPostgreSQL`, `common.UsingSQLite`, and `common.UsingMySQL` flags for DB-specific branches. - Use `common.UsingMainDatabase(...)` for primary database branches and `common.UsingLogDatabase(...)` for log database branches.
- Do not use database-specific features without cross-DB fallback, including MySQL-only functions, PostgreSQL-only operators, SQLite-unsupported `ALTER COLUMN`, or database-specific JSON column types without a `TEXT` fallback. - Do not use database-specific features without cross-DB fallback, including MySQL-only functions, PostgreSQL-only operators, SQLite-unsupported `ALTER COLUMN`, or database-specific JSON column types without a `TEXT` fallback.
- Migrations must work on all three databases. For SQLite, use `ALTER TABLE ... ADD COLUMN` instead of `ALTER COLUMN` (see `model/main.go` for patterns). - Migrations must work on all three databases. For SQLite, use `ALTER TABLE ... ADD COLUMN` instead of `ALTER COLUMN` (see `model/main.go` for patterns).
- Avoid GORM boolean default tags such as `gorm:"default:true"` when the default is a business rule already enforced by code. MySQL and PostgreSQL can normalize boolean defaults differently, causing GORM `AutoMigrate` to repeatedly issue `ALTER TABLE` on restart. Prefer setting these defaults in request/model normalization, hooks, constructors, or service logic; do not replace `default:true` with `default:1` unless the behavior is verified across SQLite, MySQL, and PostgreSQL. - Avoid GORM boolean default tags such as `gorm:"default:true"` when the default is a business rule already enforced by code. MySQL and PostgreSQL can normalize boolean defaults differently, causing GORM `AutoMigrate` to repeatedly issue `ALTER TABLE` on restart. Prefer setting these defaults in request/model normalization, hooks, constructors, or service logic; do not replace `default:true` with `default:1` unless the behavior is verified across SQLite, MySQL, and PostgreSQL.
......
common/database.go
package common package common
type DatabaseType string
const ( const (
DatabaseTypeMySQL = "mysql" DatabaseTypeMySQL DatabaseType = "mysql"
DatabaseTypeSQLite = "sqlite" DatabaseTypeSQLite DatabaseType = "sqlite"
DatabaseTypePostgreSQL = "postgres" DatabaseTypePostgreSQL DatabaseType = "postgres"
DatabaseTypeClickHouse DatabaseType = "clickhouse"
) )
var UsingSQLite = false var mainDatabaseType = DatabaseTypeSQLite
var UsingPostgreSQL = false var logDatabaseType = DatabaseTypeSQLite
var LogSqlType = DatabaseTypeSQLite // Default to SQLite for logging SQL queries
var UsingMySQL = false func MainDatabaseType() DatabaseType {
var UsingClickHouse = false return mainDatabaseType
}
func LogDatabaseType() DatabaseType {
return logDatabaseType
}
func SetMainDatabaseType(databaseType DatabaseType) {
mainDatabaseType = databaseType
}
func SetLogDatabaseType(databaseType DatabaseType) {
logDatabaseType = databaseType
}
func SetDatabaseTypes(mainType DatabaseType, logType DatabaseType) {
mainDatabaseType = mainType
logDatabaseType = logType
}
func UsingMainDatabase(databaseType DatabaseType) bool {
return mainDatabaseType == databaseType
}
func UsingLogDatabase(databaseType DatabaseType) bool {
return logDatabaseType == databaseType
}
var SQLitePath = "one-api.db?_busy_timeout=30000" var SQLitePath = "one-api.db?_busy_timeout=30000"
common/utils.go
...@@ -2,7 +2,9 @@ package common ...@@ -2,7 +2,9 @@ package common
import ( import (
crand "crypto/rand" crand "crypto/rand"
"crypto/sha256"
"encoding/base64" "encoding/base64"
"encoding/hex"
"encoding/json" "encoding/json"
"fmt" "fmt"
"html/template" "html/template"
...@@ -15,6 +17,7 @@ import ( ...@@ -15,6 +17,7 @@ import (
"os" "os"
"os/exec" "os/exec"
"runtime" "runtime"
"runtime/debug"
"strconv" "strconv"
"strings" "strings"
"time" "time"
...@@ -264,7 +267,19 @@ func GetTimestamp() int64 { ...@@ -264,7 +267,19 @@ func GetTimestamp() int64 {
func GetTimeString() string { func GetTimeString() string {
now := time.Now().UTC() now := time.Now().UTC()
return fmt.Sprintf("%s%d", now.Format("20060102150405"), now.UnixNano()%1e9) return fmt.Sprintf("%s%09d", now.Format("20060102150405"), now.UnixNano()%1e9)
}
var requestIdPrefix = func() string {
if bi, ok := debug.ReadBuildInfo(); ok && bi.Main.Path != "" {
h := sha256.Sum256([]byte(bi.Main.Path))
return hex.EncodeToString(h[:4])
}
return GetRandomString(8)
}()
func NewRequestId() string {
return GetTimeString() + requestIdPrefix + GetRandomString(8)
} }
func Max(a int, b int) int { func Max(a int, b int) int {
......
controller/model_list_test.go
...@@ -32,9 +32,7 @@ func setupModelListControllerTestDB(t *testing.T) *gorm.DB { ...@@ -32,9 +32,7 @@ func setupModelListControllerTestDB(t *testing.T) *gorm.DB {
initModelListColumnNames(t) initModelListColumnNames(t)
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
common.UsingSQLite = true common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
common.UsingMySQL = false
common.UsingPostgreSQL = false
common.RedisEnabled = false common.RedisEnabled = false
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
...@@ -60,16 +58,13 @@ func initModelListColumnNames(t *testing.T) { ...@@ -60,16 +58,13 @@ func initModelListColumnNames(t *testing.T) {
originalIsMasterNode := common.IsMasterNode originalIsMasterNode := common.IsMasterNode
originalSQLitePath := common.SQLitePath originalSQLitePath := common.SQLitePath
originalUsingSQLite := common.UsingSQLite originalMainDatabaseType := common.MainDatabaseType()
originalUsingMySQL := common.UsingMySQL originalLogDatabaseType := common.LogDatabaseType()
originalUsingPostgreSQL := common.UsingPostgreSQL
originalSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") originalSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
defer func() { defer func() {
common.IsMasterNode = originalIsMasterNode common.IsMasterNode = originalIsMasterNode
common.SQLitePath = originalSQLitePath common.SQLitePath = originalSQLitePath
common.UsingSQLite = originalUsingSQLite common.SetDatabaseTypes(originalMainDatabaseType, originalLogDatabaseType)
common.UsingMySQL = originalUsingMySQL
common.UsingPostgreSQL = originalUsingPostgreSQL
if hadSQLDSN { if hadSQLDSN {
require.NoError(t, os.Setenv("SQL_DSN", originalSQLDSN)) require.NoError(t, os.Setenv("SQL_DSN", originalSQLDSN))
} else { } else {
...@@ -79,9 +74,7 @@ func initModelListColumnNames(t *testing.T) { ...@@ -79,9 +74,7 @@ func initModelListColumnNames(t *testing.T) {
common.IsMasterNode = false common.IsMasterNode = false
common.SQLitePath = fmt.Sprintf("file:%s_init?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) common.SQLitePath = fmt.Sprintf("file:%s_init?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
common.UsingSQLite = false common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
common.UsingMySQL = false
common.UsingPostgreSQL = false
require.NoError(t, os.Setenv("SQL_DSN", "local")) require.NoError(t, os.Setenv("SQL_DSN", "local"))
require.NoError(t, model.InitDB()) require.NoError(t, model.InitDB())
......
controller/setup.go
...@@ -36,15 +36,7 @@ func GetSetup(c *gin.Context) { ...@@ -36,15 +36,7 @@ func GetSetup(c *gin.Context) {
return return
} }
setup.RootInit = model.RootUserExists() setup.RootInit = model.RootUserExists()
if common.UsingMySQL { setup.DatabaseType = string(common.MainDatabaseType())
setup.DatabaseType = "mysql"
}
if common.UsingPostgreSQL {
setup.DatabaseType = "postgres"
}
if common.UsingSQLite {
setup.DatabaseType = "sqlite"
}
c.JSON(200, gin.H{ c.JSON(200, gin.H{
"success": true, "success": true,
"data": setup, "data": setup,
......
controller/token_test.go
...@@ -48,21 +48,21 @@ type sqliteColumnInfo struct { ...@@ -48,21 +48,21 @@ type sqliteColumnInfo struct {
} }
type legacyToken struct { type legacyToken struct {
Id int `gorm:"primaryKey"` Id int `gorm:"primaryKey"`
UserId int `gorm:"index"` UserId int `gorm:"index"`
Key string `gorm:"column:key;type:char(48);uniqueIndex"` Key string `gorm:"column:key;type:char(48);uniqueIndex"`
Status int `gorm:"default:1"` Status int `gorm:"default:1"`
Name string `gorm:"index"` Name string `gorm:"index"`
CreatedTime int64 `gorm:"bigint"` CreatedTime int64 `gorm:"bigint"`
AccessedTime int64 `gorm:"bigint"` AccessedTime int64 `gorm:"bigint"`
ExpiredTime int64 `gorm:"bigint;default:-1"` ExpiredTime int64 `gorm:"bigint;default:-1"`
RemainQuota int `gorm:"default:0"` RemainQuota int `gorm:"default:0"`
UnlimitedQuota bool UnlimitedQuota bool
ModelLimitsEnabled bool ModelLimitsEnabled bool
ModelLimits string `gorm:"type:text"` ModelLimits string `gorm:"type:text"`
AllowIps *string `gorm:"default:''"` AllowIps *string `gorm:"default:''"`
UsedQuota int `gorm:"default:0"` UsedQuota int `gorm:"default:0"`
Group string `gorm:"column:group;default:''"` Group string `gorm:"column:group;default:''"`
CrossGroupRetry bool CrossGroupRetry bool
DeletedAt gorm.DeletedAt `gorm:"index"` DeletedAt gorm.DeletedAt `gorm:"index"`
} }
...@@ -75,9 +75,7 @@ func openTokenControllerTestDB(t *testing.T) *gorm.DB { ...@@ -75,9 +75,7 @@ func openTokenControllerTestDB(t *testing.T) *gorm.DB {
t.Helper() t.Helper()
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
common.UsingSQLite = true common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
common.UsingMySQL = false
common.UsingPostgreSQL = false
common.RedisEnabled = false common.RedisEnabled = false
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
...@@ -119,22 +117,23 @@ func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*g ...@@ -119,22 +117,23 @@ func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*g
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
common.RedisEnabled = false common.RedisEnabled = false
common.UsingSQLite = false
common.UsingMySQL = dialect == "mysql"
common.UsingPostgreSQL = dialect == "postgres"
var ( var (
db *gorm.DB db *gorm.DB
err error dbType common.DatabaseType
err error
) )
switch dialect { switch dialect {
case "mysql": case "mysql":
dbType = common.DatabaseTypeMySQL
db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{}) db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{})
case "postgres": case "postgres":
dbType = common.DatabaseTypePostgreSQL
db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{}) db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
default: default:
t.Fatalf("unsupported dialect %q", dialect) t.Fatalf("unsupported dialect %q", dialect)
} }
common.SetDatabaseTypes(dbType, dbType)
if err != nil { if err != nil {
t.Fatalf("failed to open %s db: %v", dialect, err) t.Fatalf("failed to open %s db: %v", dialect, err)
} }
......
docker-compose.yml
...@@ -28,6 +28,9 @@ services: ...@@ -28,6 +28,9 @@ services:
environment: environment:
- SQL_DSN=postgresql://root:123456@postgres:5432/new-api # ⚠️ IMPORTANT: Change the password in production! - SQL_DSN=postgresql://root:123456@postgres:5432/new-api # ⚠️ IMPORTANT: Change the password in production!
# - SQL_DSN=root:123456@tcp(mysql:3306)/new-api # Point to the mysql service, uncomment if using MySQL # - SQL_DSN=root:123456@tcp(mysql:3306)/new-api # Point to the mysql service, uncomment if using MySQL
# - LOG_SQL_DSN=postgresql://root:123456@postgres:5432/new-api-log # OPTIONAL: If you want a separate database for logging, uncomment and set this
# - LOG_SQL_DSN=clickhouse://default:123456@clickhouse:9000/new_api_logs # OPTIONAL: Use ClickHouse for logs only; also uncomment clickhouse in depends_on and the clickhouse service below
# - LOG_SQL_CLICKHOUSE_TTL_DAYS=0 # OPTIONAL: ClickHouse log retention days. Unset or 0 disables automatic deletion; set to e.g. 30 to keep 30 days
- REDIS_CONN_STRING=redis://:123456@redis:6379 # ⚠️ IMPORTANT: Change the password in production! - REDIS_CONN_STRING=redis://:123456@redis:6379 # ⚠️ IMPORTANT: Change the password in production!
- TZ=Asia/Shanghai - TZ=Asia/Shanghai
- ERROR_LOG_ENABLED=true # 是否启用错误日志记录 (Whether to enable error log recording) - ERROR_LOG_ENABLED=true # 是否启用错误日志记录 (Whether to enable error log recording)
...@@ -45,6 +48,7 @@ services: ...@@ -45,6 +48,7 @@ services:
- redis - redis
- postgres - postgres
# - mysql # Uncomment if using MySQL # - mysql # Uncomment if using MySQL
# - clickhouse # Uncomment if using ClickHouse for LOG_SQL_DSN
networks: networks:
- new-api-network - new-api-network
healthcheck: healthcheck:
...@@ -90,9 +94,27 @@ services: ...@@ -90,9 +94,27 @@ services:
# ports: # ports:
# - "3306:3306" # Uncomment if you need to access MySQL from outside Docker # - "3306:3306" # Uncomment if you need to access MySQL from outside Docker
# clickhouse:
# image: clickhouse/clickhouse-server:24.8
# container_name: clickhouse
# restart: always
# environment:
# CLICKHOUSE_DB: new_api_logs
# CLICKHOUSE_USER: default
# CLICKHOUSE_PASSWORD: 123456 # ⚠️ IMPORTANT: Change this password in production!
# CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT: 1
# volumes:
# - clickhouse_data:/var/lib/clickhouse
# networks:
# - new-api-network
# ports:
# - "8123:8123" # HTTP interface, uncomment if you need external access
# - "9000:9000" # Native interface used by the LOG_SQL_DSN example above
volumes: volumes:
pg_data: pg_data:
# mysql_data: # mysql_data:
# clickhouse_data:
networks: networks:
new-api-network: new-api-network:
......
go.mod
...@@ -60,7 +60,23 @@ require ( ...@@ -60,7 +60,23 @@ require (
gorm.io/gorm v1.25.2 gorm.io/gorm v1.25.2
) )
require github.com/waffo-com/waffo-pancake-sdk-go v0.3.1 require (
github.com/waffo-com/waffo-pancake-sdk-go v0.3.1
gorm.io/driver/clickhouse v0.6.0
)
require (
github.com/ClickHouse/ch-go v0.58.2 // indirect
github.com/ClickHouse/clickhouse-go/v2 v2.15.0 // indirect
github.com/go-faster/city v1.0.1 // indirect
github.com/go-faster/errors v0.6.1 // indirect
github.com/hashicorp/go-version v1.6.0 // indirect
github.com/paulmach/orb v0.10.0 // indirect
github.com/pierrec/lz4/v4 v4.1.18 // indirect
github.com/segmentio/asm v1.2.0 // indirect
go.opentelemetry.io/otel v1.19.0 // indirect
go.opentelemetry.io/otel/trace v1.19.0 // indirect
)
require ( require (
github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/DmitriyVTitov/size v1.5.0 // indirect
......
go.sum
This source diff could not be displayed because it is too large. You can view the blob instead.
middleware/request-id.go
...@@ -2,25 +2,14 @@ package middleware ...@@ -2,25 +2,14 @@ package middleware
import ( import (
"context" "context"
"crypto/sha256"
"encoding/hex"
"runtime/debug"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
var _bp = func() string {
if bi, ok := debug.ReadBuildInfo(); ok && bi.Main.Path != "" {
h := sha256.Sum256([]byte(bi.Main.Path))
return hex.EncodeToString(h[:4])
}
return common.GetRandomString(8)
}()
func RequestId() func(c *gin.Context) { func RequestId() func(c *gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
id := common.GetTimeString() + _bp + common.GetRandomString(8) id := common.NewRequestId()
c.Set(common.RequestIdKey, id) c.Set(common.RequestIdKey, id)
ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id) ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id)
c.Request = c.Request.WithContext(ctx) c.Request = c.Request.WithContext(ctx)
......
model/ability.go
...@@ -113,7 +113,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha ...@@ -113,7 +113,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
if err != nil { if err != nil {
return nil, err 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 err = channelQuery.Order("weight DESC").Find(&abilities).Error
} else { } else {
err = channelQuery.Order("weight DESC").Find(&abilities).Error err = channelQuery.Order("weight DESC").Find(&abilities).Error
...@@ -341,7 +341,7 @@ func FixAbility() (int, int, error) { ...@@ -341,7 +341,7 @@ func FixAbility() (int, int, error) {
defer fixLock.Unlock() defer fixLock.Unlock()
// truncate abilities table // truncate abilities table
if common.UsingSQLite { if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
err := DB.Exec("DELETE FROM abilities").Error err := DB.Exec("DELETE FROM abilities").Error
if err != nil { if err != nil {
common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error())) common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
......
model/channel.go
...@@ -138,7 +138,7 @@ func NormalizeChannelGroupFilter(group string) string { ...@@ -138,7 +138,7 @@ func NormalizeChannelGroupFilter(group string) string {
} }
func channelGroupFilterCondition() string { func channelGroupFilterCondition() string {
if common.UsingMySQL { if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
return `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ? ESCAPE '!'` return `CONCAT(',', ` + commonGroupCol + `, ',') LIKE ? ESCAPE '!'`
} }
return `(',' || ` + commonGroupCol + ` || ',') LIKE ? ESCAPE '!'` return `(',' || ` + commonGroupCol + ` || ',') LIKE ? ESCAPE '!'`
...@@ -381,13 +381,13 @@ func SearchChannels(keyword string, group string, model string, idSort bool, sor ...@@ -381,13 +381,13 @@ func SearchChannels(keyword string, group string, model string, idSort bool, sor
modelsCol := "`models`" modelsCol := "`models`"
// 如果是 PostgreSQL，使用双引号 // 如果是 PostgreSQL，使用双引号
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
modelsCol = `"models"` modelsCol = `"models"`
} }
baseURLCol := "`base_url`" baseURLCol := "`base_url`"
// 如果是 PostgreSQL，使用双引号 // 如果是 PostgreSQL，使用双引号
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
baseURLCol = `"base_url"` baseURLCol = `"base_url"`
} }
...@@ -898,13 +898,13 @@ func SearchTags(keyword string, group string, model string, idSort bool) ([]*str ...@@ -898,13 +898,13 @@ func SearchTags(keyword string, group string, model string, idSort bool) ([]*str
modelsCol := "`models`" modelsCol := "`models`"
// 如果是 PostgreSQL，使用双引号 // 如果是 PostgreSQL，使用双引号
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
modelsCol = `"models"` modelsCol = `"models"`
} }
baseURLCol := "`base_url`" baseURLCol := "`base_url`"
// 如果是 PostgreSQL，使用双引号 // 如果是 PostgreSQL，使用双引号
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
baseURLCol = `"base_url"` baseURLCol = `"base_url"`
} }
......
model/checkin.go
...@@ -82,7 +82,7 @@ func UserCheckin(userId int) (*Checkin, error) { ...@@ -82,7 +82,7 @@ func UserCheckin(userId int) (*Checkin, error) {
} }
// 根据数据库类型选择不同的策略 // 根据数据库类型选择不同的策略
if common.UsingSQLite { if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
// SQLite 不支持嵌套事务，使用顺序操作 + 手动回滚 // SQLite 不支持嵌套事务，使用顺序操作 + 手动回滚
return userCheckinWithoutTransaction(checkin, userId, quotaAwarded) return userCheckinWithoutTransaction(checkin, userId, quotaAwarded)
} }
......
model/db_time.go
...@@ -8,9 +8,9 @@ func GetDBTimestamp() int64 { ...@@ -8,9 +8,9 @@ func GetDBTimestamp() int64 {
var ts int64 var ts int64
var err error var err error
switch { switch {
case common.UsingPostgreSQL: case common.UsingMainDatabase(common.DatabaseTypePostgreSQL):
err = DB.Raw("SELECT EXTRACT(EPOCH FROM NOW())::bigint").Scan(&ts).Error 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 err = DB.Raw("SELECT strftime('%s','now')").Scan(&ts).Error
default: default:
err = DB.Raw("SELECT UNIX_TIMESTAMP()").Scan(&ts).Error err = DB.Raw("SELECT UNIX_TIMESTAMP()").Scan(&ts).Error
......
model/log.go
...@@ -67,6 +67,27 @@ const ( ...@@ -67,6 +67,27 @@ const (
LogTypeLogin = 7 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) { func formatUserLogs(logs []*Log, startIdx int) {
for i := range logs { for i := range logs {
logs[i].ChannelName = "" logs[i].ChannelName = ""
...@@ -81,12 +102,16 @@ func formatUserLogs(logs []*Log, startIdx int) { ...@@ -81,12 +102,16 @@ func formatUserLogs(logs []*Log, startIdx int) {
delete(otherMap, "stream_status") delete(otherMap, "stream_status")
} }
logs[i].Other = common.MapToJsonStr(otherMap) logs[i].Other = common.MapToJsonStr(otherMap)
logs[i].Id = startIdx + i + 1
} }
assignDisplayLogIds(logs, startIdx)
} }
func GetLogByTokenId(tokenId int) (logs []*Log, err error) { 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) formatUserLogs(logs, 0)
return logs, err return logs, err
} }
...@@ -103,7 +128,7 @@ func RecordLog(userId int, logType int, content string) { ...@@ -103,7 +128,7 @@ func RecordLog(userId int, logType int, content string) {
Type: logType, Type: logType,
Content: content, Content: content,
} }
err := LOG_DB.Create(log).Error err := createLog(log)
if err != nil { if err != nil {
common.SysLog("failed to record log: " + err.Error()) common.SysLog("failed to record log: " + err.Error())
} }
...@@ -128,7 +153,7 @@ func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo m ...@@ -128,7 +153,7 @@ func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo m
} }
log.Other = common.MapToJsonStr(other) 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()) common.SysLog("failed to record log: " + err.Error())
} }
} }
...@@ -165,7 +190,7 @@ func RecordLoginLog(userId int, username string, content string, ip string, acti ...@@ -165,7 +190,7 @@ func RecordLoginLog(userId int, username string, content string, ip string, acti
Ip: ip, Ip: ip,
Other: common.MapToJsonStr(other), 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()) common.SysLog("failed to record login log: " + err.Error())
} }
} }
...@@ -196,7 +221,7 @@ func RecordOperationAuditLog(logUserId int, content string, ip string, action st ...@@ -196,7 +221,7 @@ func RecordOperationAuditLog(logUserId int, content string, ip string, action st
Ip: ip, Ip: ip,
Other: common.MapToJsonStr(other), 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()) common.SysLog("failed to record operation audit log: " + err.Error())
} }
} }
...@@ -223,7 +248,7 @@ func RecordTopupLog(userId int, content string, callerIp string, paymentMethod s ...@@ -223,7 +248,7 @@ func RecordTopupLog(userId int, content string, callerIp string, paymentMethod s
Ip: callerIp, Ip: callerIp,
Other: common.MapToJsonStr(other), Other: common.MapToJsonStr(other),
} }
err := LOG_DB.Create(log).Error err := createLog(log)
if err != nil { if err != nil {
common.SysLog("failed to record topup log: " + err.Error()) common.SysLog("failed to record topup log: " + err.Error())
} }
...@@ -269,7 +294,7 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, ...@@ -269,7 +294,7 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string,
UpstreamRequestId: upstreamRequestId, UpstreamRequestId: upstreamRequestId,
Other: otherStr, Other: otherStr,
} }
err := LOG_DB.Create(log).Error err := createLog(log)
if err != nil { if err != nil {
logger.LogError(c, "failed to record log: "+err.Error()) logger.LogError(c, "failed to record log: "+err.Error())
} }
...@@ -333,7 +358,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) ...@@ -333,7 +358,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
UpstreamRequestId: upstreamRequestId, UpstreamRequestId: upstreamRequestId,
Other: otherStr, Other: otherStr,
} }
err := LOG_DB.Create(log).Error err := createLog(log)
if err != nil { if err != nil {
logger.LogError(c, "failed to record log: "+err.Error()) logger.LogError(c, "failed to record log: "+err.Error())
} }
...@@ -393,7 +418,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { ...@@ -393,7 +418,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
Group: params.Group, Group: params.Group,
Other: common.MapToJsonStr(params.Other), Other: common.MapToJsonStr(params.Other),
} }
err := LOG_DB.Create(log).Error err := createLog(log)
if err != nil { if err != nil {
common.SysLog("failed to record task billing log: " + err.Error()) common.SysLog("failed to record task billing log: " + err.Error())
} }
...@@ -453,10 +478,17 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName ...@@ -453,10 +478,17 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
if err != nil { if err != nil {
return nil, 0, err 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 { if err != nil {
return nil, 0, err return nil, 0, err
} }
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
assignDisplayLogIds(logs, startIdx)
}
channelIds := types.NewSet[int]() channelIds := types.NewSet[int]()
for _, log := range logs { for _, log := range logs {
...@@ -537,7 +569,11 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int ...@@ -537,7 +569,11 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int
common.SysError("failed to count user logs: " + err.Error()) common.SysError("failed to count user logs: " + err.Error())
return nil, 0, errors.New("查询日志失败") 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 { if err != nil {
common.SysError("failed to search user logs: " + err.Error()) common.SysError("failed to search user logs: " + err.Error())
return nil, 0, errors.New("查询日志失败") return nil, 0, errors.New("查询日志失败")
...@@ -554,10 +590,10 @@ type Stat struct { ...@@ -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) { 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创建单独的查询 // 为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 { if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil {
return stat, err return stat, err
...@@ -610,7 +646,7 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa ...@@ -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) { 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 != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)
} }
...@@ -631,6 +667,54 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa ...@@ -631,6 +667,54 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa
} }
func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) { 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 var total int64 = 0
for { for {
......
model/main.go
...@@ -3,6 +3,7 @@ package model ...@@ -3,6 +3,7 @@ package model
import ( import (
"fmt" "fmt"
"log" "log"
"net/url"
"os" "os"
"strings" "strings"
"sync" "sync"
...@@ -12,6 +13,7 @@ import ( ...@@ -12,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/glebarez/sqlite" "github.com/glebarez/sqlite"
"gorm.io/driver/clickhouse"
"gorm.io/driver/mysql" "gorm.io/driver/mysql"
"gorm.io/driver/postgres" "gorm.io/driver/postgres"
"gorm.io/gorm" "gorm.io/gorm"
...@@ -27,7 +29,7 @@ var logGroupCol string ...@@ -27,7 +29,7 @@ var logGroupCol string
func initCol() { func initCol() {
// init common column names // init common column names
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
commonGroupCol = `"group"` commonGroupCol = `"group"`
commonKeyCol = `"key"` commonKeyCol = `"key"`
commonTrueVal = "true" commonTrueVal = "true"
...@@ -38,27 +40,14 @@ func initCol() { ...@@ -38,27 +40,14 @@ func initCol() {
commonTrueVal = "1" commonTrueVal = "1"
commonFalseVal = "0" commonFalseVal = "0"
} }
if os.Getenv("LOG_SQL_DSN") != "" { switch common.LogDatabaseType() {
switch common.LogSqlType { case common.DatabaseTypePostgreSQL:
case common.DatabaseTypePostgreSQL: logGroupCol = `"group"`
logGroupCol = `"group"` logKeyCol = `"key"`
logKeyCol = `"key"` default:
default: logGroupCol = "`group`"
logGroupCol = commonGroupCol logKeyCol = "`key`"
logKeyCol = commonKeyCol
}
} else {
// LOG_SQL_DSN 为空时，日志数据库与主数据库相同
if common.UsingPostgreSQL {
logGroupCol = `"group"`
logKeyCol = `"key"`
} else {
logGroupCol = commonGroupCol
logKeyCol = commonKeyCol
}
} }
// log sql type and database type
//common.SysLog("Using Log SQL Type: " + common.LogSqlType)
} }
var DB *gorm.DB var DB *gorm.DB
...@@ -115,37 +104,56 @@ func CheckSetup() { ...@@ -115,37 +104,56 @@ func CheckSetup() {
} }
} }
func chooseDB(envName string, isLog bool) (*gorm.DB, error) { func isClickHouseDSN(dsn string) bool {
defer func() { return strings.HasPrefix(dsn, "clickhouse://") ||
initCol() 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) dsn := os.Getenv(envName)
if dsn != "" { 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://") { if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
// Use PostgreSQL // Use PostgreSQL
common.SysLog("using PostgreSQL as database") common.SysLog("using PostgreSQL as database")
if !isLog { db, err := gorm.Open(postgres.New(postgres.Config{
common.UsingPostgreSQL = true
} else {
common.LogSqlType = common.DatabaseTypePostgreSQL
}
return gorm.Open(postgres.New(postgres.Config{
DSN: dsn, DSN: dsn,
PreferSimpleProtocol: true, // disables implicit prepared statement usage PreferSimpleProtocol: true, // disables implicit prepared statement usage
}), &gorm.Config{ }), &gorm.Config{
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
return db, common.DatabaseTypePostgreSQL, err
} }
if strings.HasPrefix(dsn, "local") { if strings.HasPrefix(dsn, "local") {
common.SysLog("SQL_DSN not set, using SQLite as database") common.SysLog("SQL_DSN not set, using SQLite as database")
if !isLog { db, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
common.UsingSQLite = true
} else {
common.LogSqlType = common.DatabaseTypeSQLite
}
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
return db, common.DatabaseTypeSQLite, err
} }
// Use MySQL // Use MySQL
common.SysLog("using MySQL as database") common.SysLog("using MySQL as database")
...@@ -157,32 +165,33 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) { ...@@ -157,32 +165,33 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) {
dsn += "?parseTime=true" dsn += "?parseTime=true"
} }
} }
if !isLog { db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
common.UsingMySQL = true
} else {
common.LogSqlType = common.DatabaseTypeMySQL
}
return gorm.Open(mysql.Open(dsn), &gorm.Config{
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
return db, common.DatabaseTypeMySQL, err
} }
// Use SQLite // Use SQLite
common.SysLog("SQL_DSN not set, using SQLite as database") common.SysLog("SQL_DSN not set, using SQLite as database")
common.UsingSQLite = true db, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
PrepareStmt: true, // precompile SQL PrepareStmt: true, // precompile SQL
}) })
return db, common.DatabaseTypeSQLite, err
} }
func InitDB() (err error) { func InitDB() (err error) {
db, err := chooseDB("SQL_DSN", false) db, dbType, err := chooseDB("SQL_DSN", false)
if err == nil { if err == nil {
common.SetMainDatabaseType(dbType)
if os.Getenv("LOG_SQL_DSN") == "" {
common.SetLogDatabaseType(dbType)
}
initCol()
if common.DebugEnabled { if common.DebugEnabled {
db = db.Debug() db = db.Debug()
} }
DB = db DB = db
// MySQL charset/collation startup check: ensure Chinese-capable charset // MySQL charset/collation startup check: ensure Chinese-capable charset
if common.UsingMySQL { if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
if err := checkMySQLChineseSupport(DB); err != nil { if err := checkMySQLChineseSupport(DB); err != nil {
panic(err) panic(err)
} }
...@@ -198,7 +207,7 @@ func InitDB() (err error) { ...@@ -198,7 +207,7 @@ func InitDB() (err error) {
if !common.IsMasterNode { if !common.IsMasterNode {
return nil 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 //_, _ = sqlDB.Exec("ALTER TABLE channels MODIFY model_mapping TEXT;") // TODO: delete this line when most users have upgraded
} }
common.SysLog("database migration started") common.SysLog("database migration started")
...@@ -213,16 +222,20 @@ func InitDB() (err error) { ...@@ -213,16 +222,20 @@ func InitDB() (err error) {
func InitLogDB() (err error) { func InitLogDB() (err error) {
if os.Getenv("LOG_SQL_DSN") == "" { if os.Getenv("LOG_SQL_DSN") == "" {
LOG_DB = DB LOG_DB = DB
common.SetLogDatabaseType(common.MainDatabaseType())
initCol()
return return
} }
db, err := chooseDB("LOG_SQL_DSN", true) db, dbType, err := chooseDB("LOG_SQL_DSN", true)
if err == nil { if err == nil {
common.SetLogDatabaseType(dbType)
initCol()
if common.DebugEnabled { if common.DebugEnabled {
db = db.Debug() db = db.Debug()
} }
LOG_DB = db LOG_DB = db
// If log DB is MySQL, also ensure Chinese-capable charset // 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 { if err := checkMySQLChineseSupport(LOG_DB); err != nil {
panic(err) panic(err)
} }
...@@ -285,7 +298,7 @@ func migrateDB() error { ...@@ -285,7 +298,7 @@ func migrateDB() error {
if err != nil { if err != nil {
return err return err
} }
if common.UsingSQLite { if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
if err := ensureSubscriptionPlanTableSQLite(); err != nil { if err := ensureSubscriptionPlanTableSQLite(); err != nil {
return err return err
} }
...@@ -354,7 +367,7 @@ func migrateDBFast() error { ...@@ -354,7 +367,7 @@ func migrateDBFast() error {
return err return err
} }
} }
if common.UsingSQLite { if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
if err := ensureSubscriptionPlanTableSQLite(); err != nil { if err := ensureSubscriptionPlanTableSQLite(); err != nil {
return err return err
} }
...@@ -368,11 +381,99 @@ func migrateDBFast() error { ...@@ -368,11 +381,99 @@ func migrateDBFast() error {
} }
func migrateLOGDB() error { func migrateLOGDB() error {
var err error if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
if err = LOG_DB.AutoMigrate(&Log{}); err != nil { 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 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 { type sqliteColumnDef struct {
...@@ -381,7 +482,7 @@ type sqliteColumnDef struct { ...@@ -381,7 +482,7 @@ type sqliteColumnDef struct {
} }
func ensureSubscriptionPlanTableSQLite() error { func ensureSubscriptionPlanTableSQLite() error {
if !common.UsingSQLite { if !common.UsingMainDatabase(common.DatabaseTypeSQLite) {
return nil return nil
} }
tableName := "subscription_plans" tableName := "subscription_plans"
...@@ -463,7 +564,7 @@ PRIMARY KEY (` + "`id`" + `) ...@@ -463,7 +564,7 @@ PRIMARY KEY (` + "`id`" + `)
// This is safe to run multiple times - it checks the column type first // This is safe to run multiple times - it checks the column type first
func migrateTokenModelLimitsToText() error { func migrateTokenModelLimitsToText() error {
// SQLite uses type affinity, so TEXT and VARCHAR are effectively the same — no migration needed // 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 return nil
} }
...@@ -479,7 +580,7 @@ func migrateTokenModelLimitsToText() error { ...@@ -479,7 +580,7 @@ func migrateTokenModelLimitsToText() error {
} }
var alterSQL string var alterSQL string
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
var dataType string var dataType string
if err := DB.Raw(`SELECT data_type FROM information_schema.columns if err := DB.Raw(`SELECT data_type FROM information_schema.columns
WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`, WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`,
...@@ -489,7 +590,7 @@ func migrateTokenModelLimitsToText() error { ...@@ -489,7 +590,7 @@ func migrateTokenModelLimitsToText() error {
return nil return nil
} }
alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE text`, tableName, columnName) 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 var columnType string
if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
...@@ -517,7 +618,7 @@ func migrateTokenModelLimitsToText() error { ...@@ -517,7 +618,7 @@ func migrateTokenModelLimitsToText() error {
func migrateSubscriptionPlanPriceAmount() { func migrateSubscriptionPlanPriceAmount() {
// SQLite doesn't support ALTER COLUMN, and its type affinity handles this automatically // 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 // Skip early to avoid GORM parsing the existing table DDL which may cause issues
if common.UsingSQLite { if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
return return
} }
...@@ -535,7 +636,7 @@ func migrateSubscriptionPlanPriceAmount() { ...@@ -535,7 +636,7 @@ func migrateSubscriptionPlanPriceAmount() {
} }
var alterSQL string var alterSQL string
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
// PostgreSQL: Check if already decimal/numeric // PostgreSQL: Check if already decimal/numeric
var dataType string var dataType string
if err := DB.Raw(`SELECT data_type FROM information_schema.columns if err := DB.Raw(`SELECT data_type FROM information_schema.columns
...@@ -547,7 +648,7 @@ func migrateSubscriptionPlanPriceAmount() { ...@@ -547,7 +648,7 @@ func migrateSubscriptionPlanPriceAmount() {
} }
alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE decimal(10,6) USING %s::decimal(10,6)`, alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE decimal(10,6) USING %s::decimal(10,6)`,
tableName, columnName, columnName) tableName, columnName, columnName)
} else if common.UsingMySQL { } else if common.UsingMainDatabase(common.DatabaseTypeMySQL) {
// MySQL: Check if already decimal // MySQL: Check if already decimal
var columnType string var columnType string
if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
......
model/redemption.go
...@@ -122,7 +122,7 @@ func Redeem(key string, userId int) (quota int, err error) { ...@@ -122,7 +122,7 @@ func Redeem(key string, userId int) (quota int, err error) {
redemption := &Redemption{} redemption := &Redemption{}
keyCol := "`key`" keyCol := "`key`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
keyCol = `"key"` keyCol = `"key"`
} }
common.RandomSleep() common.RandomSleep()
......
model/subscription.go
...@@ -555,7 +555,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP ...@@ -555,7 +555,7 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP
return errors.New("tradeNo is empty") return errors.New("tradeNo is empty")
} }
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
var logUserId int var logUserId int
...@@ -663,7 +663,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err ...@@ -663,7 +663,7 @@ func ExpireSubscriptionOrder(tradeNo string, expectedPaymentProvider string) err
return errors.New("tradeNo is empty") return errors.New("tradeNo is empty")
} }
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
return DB.Transaction(func(tx *gorm.DB) error { return DB.Transaction(func(tx *gorm.DB) error {
......
model/task_cas_test.go
...@@ -22,7 +22,7 @@ func TestMain(m *testing.M) { ...@@ -22,7 +22,7 @@ func TestMain(m *testing.M) {
DB = db DB = db
LOG_DB = db LOG_DB = db
common.UsingSQLite = true common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
common.RedisEnabled = false common.RedisEnabled = false
common.BatchUpdateEnabled = false common.BatchUpdateEnabled = false
common.LogConsumeEnabled = true common.LogConsumeEnabled = true
......
model/topup.go
...@@ -85,7 +85,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta ...@@ -85,7 +85,7 @@ func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, ta
} }
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
...@@ -115,7 +115,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error ...@@ -115,7 +115,7 @@ func Recharge(referenceId string, customerId string, callerIp string) (err error
topUp := &TopUp{} topUp := &TopUp{}
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
...@@ -323,7 +323,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error { ...@@ -323,7 +323,7 @@ func ManualCompleteTopUp(tradeNo string, callerIp string) error {
} }
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
...@@ -398,7 +398,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string ...@@ -398,7 +398,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
topUp := &TopUp{} topUp := &TopUp{}
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
...@@ -473,7 +473,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) { ...@@ -473,7 +473,7 @@ func RechargeWaffo(tradeNo string, callerIp string) (err error) {
topUp := &TopUp{} topUp := &TopUp{}
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
...@@ -536,7 +536,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) { ...@@ -536,7 +536,7 @@ func RechargeWaffoPancake(tradeNo string) (err error) {
topUp := &TopUp{} topUp := &TopUp{}
refCol := "`trade_no`" refCol := "`trade_no`"
if common.UsingPostgreSQL { if common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
refCol = `"trade_no"` refCol = `"trade_no"`
} }
......
model/usedata_rankings.go
...@@ -49,7 +49,7 @@ func GetRankingQuotaBuckets(startTime int64, endTime int64, bucketSize int64) ([ ...@@ -49,7 +49,7 @@ func GetRankingQuotaBuckets(startTime int64, endTime int64, bucketSize int64) ([
} }
func rankingBucketExpr(bucketSize int64) string { 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("FLOOR(created_at / %d) * %d", bucketSize, bucketSize)
} }
return fmt.Sprintf("(created_at / %d) * %d", bucketSize, bucketSize) return fmt.Sprintf("(created_at / %d) * %d", bucketSize, bucketSize)
......
relay/common/relay_info.go
...@@ -459,7 +459,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { ...@@ -459,7 +459,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
reqId := common.GetContextKeyString(c, common.RequestIdKey) reqId := common.GetContextKeyString(c, common.RequestIdKey)
if reqId == "" { if reqId == "" {
reqId = common.GetTimeString() + common.GetRandomString(8) reqId = common.NewRequestId()
} }
info := &RelayInfo{ info := &RelayInfo{
Request: request, Request: request,
......
service/task_billing_test.go
...@@ -31,7 +31,7 @@ func TestMain(m *testing.M) { ...@@ -31,7 +31,7 @@ func TestMain(m *testing.M) {
model.DB = db model.DB = db
model.LOG_DB = db model.LOG_DB = db
common.UsingSQLite = true common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite)
common.RedisEnabled = false common.RedisEnabled = false
common.BatchUpdateEnabled = false common.BatchUpdateEnabled = false
common.LogConsumeEnabled = true common.LogConsumeEnabled = true
......
service/waffo_pancake_test.go deleted 100644 → 0
package service
import (
"context"
"fmt"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func setupWaffoPancakeTestDB(t *testing.T) *gorm.DB {
t.Helper()
common.UsingSQLite = true
common.UsingMySQL = false
common.UsingPostgreSQL = false
common.RedisEnabled = false
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
require.NoError(t, err)
model.DB = db
model.LOG_DB = db
require.NoError(t, db.AutoMigrate(&model.User{}, &model.TopUp{}, &model.SubscriptionOrder{}))
t.Cleanup(func() {
sqlDB, err := db.DB()
if err == nil {
_ = sqlDB.Close()
}
})
return db
}
func TestCreateWaffoPancakeCheckoutSession_RequiresOrderMerchantExternalID(t *testing.T) {
session, err := CreateWaffoPancakeCheckoutSession(context.Background(), &WaffoPancakeCreateSessionParams{
ProductID: "PROD_checkout_guard",
BuyerIdentity: WaffoPancakeBuyerIdentityFromUserID(1),
})
require.Error(t, err)
require.Nil(t, session)
require.Contains(t, err.Error(), "missing order merchant external id")
}
func TestResolveWaffoPancakeTradeNo_UsesWebhookOrderIDWhenLocalOrderExists(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
topUp := &model.TopUp{
UserId: 1,
Amount: 10,
Money: 29,
TradeNo: "ORD_5dXBtmF2HLlHfbPNm0Wcnz",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(topUp).Error)
tradeNo, err := ResolveWaffoPancakeTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "ORD_5dXBtmF2HLlHfbPNm0Wcnz",
MerchantProvidedBuyerIdentity: WaffoPancakeBuyerIdentityFromUserID(topUp.UserId),
},
})
require.NoError(t, err)
require.Equal(t, "ORD_5dXBtmF2HLlHfbPNm0Wcnz", tradeNo)
}
func TestResolveWaffoPancakeTradeNo_RejectsBuyerIdentityMismatch(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
topUp := &model.TopUp{
UserId: 42,
Amount: 10,
Money: 29,
TradeNo: "ORD_identity_mismatch_case",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(topUp).Error)
// Webhook reports the right order but a different buyer — could be a
// crossed-wires bug or a tampered payload. Either way: reject.
tradeNo, err := ResolveWaffoPancakeTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "ORD_identity_mismatch_case",
MerchantProvidedBuyerIdentity: WaffoPancakeBuyerIdentityFromUserID(99), // wrong user
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
require.Contains(t, err.Error(), "buyer identity mismatch")
}
func TestResolveWaffoPancakeTradeNo_RejectsMissingBuyerIdentity(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
topUp := &model.TopUp{
UserId: 7,
Amount: 10,
Money: 29,
TradeNo: "ORD_missing_identity",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(topUp).Error)
// An empty MerchantProvidedBuyerIdentity means the order was either created
// via the (now-deprecated) anonymous flow or the field was stripped — also
// reject so that we never credit anonymous orders to a specific user.
tradeNo, err := ResolveWaffoPancakeTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "ORD_missing_identity",
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
require.Contains(t, err.Error(), "buyer identity mismatch")
}
func TestResolveWaffoPancakeTradeNo_FailsWhenWebhookOrderIDIsUnknown(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
user := &model.User{
Id: 42,
Email: "buyer@example.com",
Username: "buyer",
Status: common.UserStatusEnabled,
}
require.NoError(t, db.Create(user).Error)
topUp := &model.TopUp{
UserId: user.Id,
Amount: 10,
Money: 29,
TradeNo: "WAFFO_PANCAKE-42-123456-abc123",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(topUp).Error)
tradeNo, err := ResolveWaffoPancakeTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "WAFFO_PANCAKE-unknown",
BuyerEmail: user.Email,
Amount: "29.00",
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
}
// Parity tests for ResolveWaffoPancakeSubscriptionTradeNo — same four cases
// as the TopUp resolver above, exercised against SubscriptionOrder records.
// Drift between the two webhook flows is a real risk because they share
// the same buyer-identity defence-in-depth pattern.
func TestResolveWaffoPancakeSubscriptionTradeNo_UsesWebhookOrderIDWhenLocalOrderExists(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
order := &model.SubscriptionOrder{
UserId: 1,
PlanId: 5,
Money: 29,
TradeNo: "WAFFO_PANCAKE_SUB-1-1700000000-abc123",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(order).Error)
tradeNo, err := ResolveWaffoPancakeSubscriptionTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "WAFFO_PANCAKE_SUB-1-1700000000-abc123",
MerchantProvidedBuyerIdentity: WaffoPancakeBuyerIdentityFromUserID(order.UserId),
},
})
require.NoError(t, err)
require.Equal(t, "WAFFO_PANCAKE_SUB-1-1700000000-abc123", tradeNo)
}
func TestResolveWaffoPancakeSubscriptionTradeNo_RejectsBuyerIdentityMismatch(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
order := &model.SubscriptionOrder{
UserId: 42,
PlanId: 5,
Money: 29,
TradeNo: "WAFFO_PANCAKE_SUB-42-mismatch",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(order).Error)
tradeNo, err := ResolveWaffoPancakeSubscriptionTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderID: "ORD_internal_pancake_id",
OrderMerchantExternalID: "WAFFO_PANCAKE_SUB-42-mismatch",
MerchantProvidedBuyerIdentity: WaffoPancakeBuyerIdentityFromUserID(99), // wrong user
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
require.Contains(t, err.Error(), "buyer identity mismatch")
}
func TestResolveWaffoPancakeSubscriptionTradeNo_RejectsMissingBuyerIdentity(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
order := &model.SubscriptionOrder{
UserId: 7,
PlanId: 5,
Money: 29,
TradeNo: "WAFFO_PANCAKE_SUB-7-missing-identity",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(order).Error)
tradeNo, err := ResolveWaffoPancakeSubscriptionTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderMerchantExternalID: "WAFFO_PANCAKE_SUB-7-missing-identity",
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
require.Contains(t, err.Error(), "buyer identity mismatch")
}
func TestResolveWaffoPancakeSubscriptionTradeNo_FailsWhenWebhookOrderIDIsUnknown(t *testing.T) {
db := setupWaffoPancakeTestDB(t)
order := &model.SubscriptionOrder{
UserId: 42,
PlanId: 5,
Money: 29,
TradeNo: "WAFFO_PANCAKE_SUB-42-real-order",
PaymentMethod: model.PaymentMethodWaffoPancake,
PaymentProvider: model.PaymentProviderWaffoPancake,
CreateTime: time.Now().Unix(),
Status: common.TopUpStatusPending,
}
require.NoError(t, db.Create(order).Error)
tradeNo, err := ResolveWaffoPancakeSubscriptionTradeNo(&WaffoPancakeWebhookEvent{
Data: WaffoPancakeWebhookData{
OrderMerchantExternalID: "WAFFO_PANCAKE_SUB-unknown",
},
})
require.Error(t, err)
require.Empty(t, tradeNo)
}
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