Files
gpt-load/internal/db/database.go
2025-07-09 09:46:38 +08:00

57 lines
1.3 KiB
Go

package db
import (
"fmt"
"gpt-load/internal/types"
"log"
"os"
"time"
"gorm.io/driver/mysql"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
var DB *gorm.DB
func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
dbConfig := configManager.GetDatabaseConfig()
if dbConfig.DSN == "" {
return nil, fmt.Errorf("DATABASE_DSN is not configured")
}
var newLogger logger.Interface
if configManager.GetLogConfig().Level == "debug" {
newLogger = logger.New(
log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
logger.Config{
SlowThreshold: time.Second, // Slow SQL threshold
LogLevel: logger.Info, // Log level
IgnoreRecordNotFoundError: true, // Ignore ErrRecordNotFound error for logger
Colorful: true, // Disable color
},
)
}
var err error
DB, err = gorm.Open(mysql.Open(dbConfig.DSN), &gorm.Config{
Logger: newLogger,
})
if err != nil {
return nil, fmt.Errorf("failed to connect to database: %w", err)
}
sqlDB, err := DB.DB()
if err != nil {
return nil, fmt.Errorf("failed to get sql.DB: %w", err)
}
// Set connection pool parameters
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(100)
sqlDB.SetConnMaxLifetime(time.Hour)
fmt.Println("Database connection initialized.")
return DB, nil
}