fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #22 from tbphp/feat-db-sqlite

Browse Source 
feat: Multi-DB Driver
This commit is contained in:
tbphp
2025-07-17 18:36:21 +08:00
committed by GitHub
parent 9c5933289a a9de834e03
commit c286e5f162
14 changed files with 253 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/config/manager.go
View File

@@ -95,10 +95,10 @@ func (m *Manager) ReloadConfig() error {
Level: utils.GetEnvOrDefault("LOG_LEVEL", "info"),
Format: utils.GetEnvOrDefault("LOG_FORMAT", "text"),
EnableFile: utils.ParseBoolean(os.Getenv("LOG_ENABLE_FILE"), false),
FilePath: utils.GetEnvOrDefault("LOG_FILE_PATH", "logs/app.log"),
FilePath: utils.GetEnvOrDefault("LOG_FILE_PATH", "./data/logs/app.log"),
},
Database: types.DatabaseConfig{
DSN: os.Getenv("DATABASE_DSN"),
DSN: utils.GetEnvOrDefault("DATABASE_DSN", "./data/gpt-load.db"),
},
RedisDSN: os.Getenv("REDIS_DSN"),
}

41
internal/db/database.go
View File

@@ -5,9 +5,13 @@ import (
"gpt-load/internal/types"
"log"
"os"
"path/filepath"
"strings"
"time"
"github.com/glebarez/sqlite"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
@@ -16,7 +20,8 @@ var DB *gorm.DB
func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
dbConfig := configManager.GetDatabaseConfig()
if dbConfig.DSN == "" {
dsn := dbConfig.DSN
if dsn == "" {
return nil, fmt.Errorf("DATABASE_DSN is not configured")
}
@@ -33,9 +38,32 @@ func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
)
}
var dialector gorm.Dialector
if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
dialector = postgres.New(postgres.Config{
DSN: dsn,
PreferSimpleProtocol: true,
})
} else if strings.Contains(dsn, "@tcp") {
if !strings.Contains(dsn, "parseTime") {
if strings.Contains(dsn, "?") {
dsn += "&parseTime=true"
} else {
dsn += "?parseTime=true"
}
}
dialector = mysql.Open(dsn)
} else {
if err := os.MkdirAll(filepath.Dir(dsn), 0755); err != nil {
return nil, fmt.Errorf("failed to create database directory: %w", err)
}
dialector = sqlite.Open(dsn + "?_busy_timeout=5000")
}
var err error
DB, err = gorm.Open(mysql.Open(dbConfig.DSN), &gorm.Config{
Logger: newLogger,
DB, err = gorm.Open(dialector, &gorm.Config{
Logger: newLogger,
PrepareStmt: true,
})
if err != nil {
return nil, fmt.Errorf("failed to connect to database: %w", err)
@@ -45,10 +73,9 @@ func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
if err != nil {
return nil, fmt.Errorf("failed to get sql.DB: %w", err)
}
// Set connection pool parameters
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(100)
// Set connection pool parameters for all drivers
sqlDB.SetMaxIdleConns(50)
sqlDB.SetMaxOpenConns(500)
sqlDB.SetConnMaxLifetime(time.Hour)
return DB, nil

24
internal/errors/errors.go
View File

@@ -3,8 +3,10 @@ package errors
import (
"errors"
"net/http"
"strings"
"github.com/go-sql-driver/mysql"
"github.com/jackc/pgx/v5/pgconn"
"gorm.io/gorm"
)
@@ -62,20 +64,28 @@ func ParseDBError(err error) *APIError {
return nil
}
// Handle record not found error
if errors.Is(err, gorm.ErrRecordNotFound) {
return ErrResourceNotFound
}
// Handle MySQL specific errors
var mysqlErr *mysql.MySQLError
if errors.As(err, &mysqlErr) {
switch mysqlErr.Number {
case 1062: // Duplicate entry for unique key
var pgErr *pgconn.PgError
if errors.As(err, &pgErr) {
if pgErr.Code == "23505" { // unique_violation
return ErrDuplicateResource
}
}
// Default to a generic database error
var mysqlErr *mysql.MySQLError
if errors.As(err, &mysqlErr) {
if mysqlErr.Number == 1062 { // Duplicate entry
return ErrDuplicateResource
}
}
// Generic check for SQLite
if strings.Contains(strings.ToLower(err.Error()), "unique constraint failed") {
return ErrDuplicateResource
}
return ErrDatabase
}

4
internal/models/types.go
View File

@@ -74,7 +74,7 @@ type APIKey struct {
// RequestLog 对应 request_logs 表
type RequestLog struct {
ID string `gorm:"type:varchar(36);primaryKey" json:"id"`
Timestamp time.Time `gorm:"type:datetime(3);not null;index" json:"timestamp"`
Timestamp time.Time `gorm:"not null;index" json:"timestamp"`
GroupID uint `gorm:"not null;index" json:"group_id"`
GroupName string `gorm:"type:varchar(255);index" json:"group_name"`
KeyValue string `gorm:"type:varchar(512)" json:"key_value"`
@@ -123,7 +123,7 @@ type ChartData struct {
// GroupHourlyStat 对应 group_hourly_stats 表，用于存储每个分组每小时的请求统计
type GroupHourlyStat struct {
ID uint `gorm:"primaryKey;autoIncrement" json:"id"`
Time time.Time `gorm:"type:datetime;not null;uniqueIndex:idx_group_time" json:"time"` // 整点时间
Time time.Time `gorm:"not null;uniqueIndex:idx_group_time" json:"time"` // 整点时间
GroupID uint `gorm:"not null;uniqueIndex:idx_group_time" json:"group_id"`
SuccessCount int64 `gorm:"not null;default:0" json:"success_count"`
FailureCount int64 `gorm:"not null;default:0" json:"failure_count"`

4
internal/services/request_log_service.go
View File

@@ -259,8 +259,8 @@ func (s *RequestLogService) writeLogsToDB(logs []*models.RequestLog) error {
err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "time"}, {Name: "group_id"}},
DoUpdates: clause.Assignments(map[string]any{
"success_count": gorm.Expr("success_count + ?", counts.Success),
"failure_count": gorm.Expr("failure_count + ?", counts.Failure),
"success_count": gorm.Expr("group_hourly_stats.success_count + ?", counts.Success),
"failure_count": gorm.Expr("group_hourly_stats.failure_count + ?", counts.Failure),
"updated_at": time.Now(),
}),
}).Create(&models.GroupHourlyStat{

2
internal/types/types.go
View File

@@ -20,7 +20,7 @@ type SystemSettings struct {
// 基础参数
AppUrl string `json:"app_url" default:"http://localhost:3001" name:"项目地址" category:"基础参数" desc:"项目的基础 URL，用于拼接分组终端节点地址。系统配置优先于环境变量 APP_URL。"`
RequestLogRetentionDays int `json:"request_log_retention_days" default:"7" name:"日志保留时长（天）" category:"基础参数" desc:"请求日志在数据库中的保留天数，0为不清理日志。" validate:"min=0"`
RequestLogWriteIntervalMinutes int `json:"request_log_write_interval_minutes" default:"5" name:"日志延迟写入周期（分钟）" category:"基础参数" desc:"请求日志从缓存写入数据库的周期（分钟），0为实时写入数据。" validate:"min=0"`
RequestLogWriteIntervalMinutes int `json:"request_log_write_interval_minutes" default:"1" name:"日志延迟写入周期（分钟）" category:"基础参数" desc:"请求日志从缓存写入数据库的周期（分钟），0为实时写入数据。" validate:"min=0"`
// 请求设置
RequestTimeout int `json:"request_timeout" default:"600" name:"请求超时（秒）" category:"请求设置" desc:"转发请求的完整生命周期超时（秒）等。" validate:"min=1"`

1
internal/utils/logger_utils.go
View File

@@ -28,7 +28,6 @@ func SetupLogger(configManager types.ConfigManager) {
})
} else {
logrus.SetFormatter(&logrus.TextFormatter{
ForceColors: true,
FullTimestamp: true,
TimestampFormat: "2006-01-02 15:04:05",
})