118 lines
2.3 KiB
Go
118 lines
2.3 KiB
Go
package config
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
// Config 应用配置
|
|
type Config struct {
|
|
DB *gorm.DB
|
|
Redis *redis.Client
|
|
}
|
|
|
|
// New 创建配置实例
|
|
func New() *Config {
|
|
return &Config{}
|
|
}
|
|
|
|
// InitDB 初始化数据库连接
|
|
func (c *Config) InitDB() error {
|
|
dsn := fmt.Sprintf("host=%s user=%s password=%s dbname=%s port=%s sslmode=%s",
|
|
getEnv("DB_HOST", "localhost"),
|
|
getEnv("DB_USER", "postgres"),
|
|
getEnv("DB_PASSWORD", "postgres"),
|
|
getEnv("DB_NAME", "appdb"),
|
|
getEnv("DB_PORT", "5432"),
|
|
getEnv("DB_SSL_MODE", "disable"),
|
|
)
|
|
|
|
// 配置 GORM
|
|
config := &gorm.Config{
|
|
Logger: logger.Default.LogMode(logger.Info),
|
|
}
|
|
|
|
if os.Getenv("APP_ENV") == "production" {
|
|
config.Logger = logger.Default.LogMode(logger.Silent)
|
|
}
|
|
|
|
db, err := gorm.Open(postgres.Open(dsn), config)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to connect to database: %w", err)
|
|
}
|
|
|
|
// 获取底层的 sqlDB
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get sqlDB: %w", err)
|
|
}
|
|
|
|
// 设置连接池
|
|
sqlDB.SetMaxIdleConns(10)
|
|
sqlDB.SetMaxOpenConns(100)
|
|
sqlDB.SetConnMaxLifetime(time.Hour)
|
|
|
|
c.DB = db
|
|
return nil
|
|
}
|
|
|
|
// CloseDB 关闭数据库连接
|
|
func (c *Config) CloseDB() error {
|
|
if c.DB != nil {
|
|
sqlDB, err := c.DB.DB()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return sqlDB.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// InitRedis 初始化 Redis 连接
|
|
func (c *Config) InitRedis() error {
|
|
opt := &redis.Options{
|
|
Addr: fmt.Sprintf("%s:%s",
|
|
getEnv("REDIS_HOST", "localhost"),
|
|
getEnv("REDIS_PORT", "6379"),
|
|
),
|
|
Password: getEnv("REDIS_PASSWORD", ""),
|
|
DB: 0,
|
|
}
|
|
|
|
client := redis.NewClient(opt)
|
|
|
|
// 测试连接
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
if err := client.Ping(ctx).Err(); err != nil {
|
|
return fmt.Errorf("failed to connect to redis: %w", err)
|
|
}
|
|
|
|
c.Redis = client
|
|
return nil
|
|
}
|
|
|
|
// CloseRedis 关闭 Redis 连接
|
|
func (c *Config) CloseRedis() error {
|
|
if c.Redis != nil {
|
|
return c.Redis.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// getEnv 获取环境变量，如果不存在则返回默认值
|
|
func getEnv(key, defaultValue string) string {
|
|
if value := os.Getenv(key); value != "" {
|
|
return value
|
|
}
|
|
return defaultValue
|
|
}
|