136 lines
3.6 KiB
Go
136 lines
3.6 KiB
Go
// common/db/postgres.go
|
||
package db
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"fmt"
|
||
"sync"
|
||
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
)
|
||
|
||
// 全局单例连接，避免重复创建连接（连接池全局复用）
|
||
var (
|
||
pgInstance *gorm.DB
|
||
pgOnce sync.Once
|
||
pgErr error
|
||
)
|
||
|
||
// InitPostgres 初始化PostgreSQL连接（单例模式，仅执行一次）
|
||
// 入参：选项模式的配置函数，用户可灵活配置
|
||
// 返回：原生*gorm.DB + 错误，业务层可直接使用GORM所有方法
|
||
func InitPostgres(opts ...PostgresOption) (*gorm.DB, error) {
|
||
pgOnce.Do(func() {
|
||
// 加载默认配置 + 覆盖用户自定义配置
|
||
options := defaultPostgresOptions()
|
||
for _, opt := range opts {
|
||
opt(options)
|
||
}
|
||
|
||
// 校验必传配置（账号、密码、数据库名不能为空）
|
||
if options.User == "" || options.Password == "" || options.DBName == "" {
|
||
pgErr = fmt.Errorf("postgres config error: user/password/dbname can not be empty")
|
||
return
|
||
}
|
||
|
||
// 拼接PostgreSQL DSN（数据源名称），GORM官方规范
|
||
dsn := fmt.Sprintf(
|
||
"host=%s port=%s user=%s password=%s dbname=%s sslmode=%s TimeZone=%s",
|
||
options.Host,
|
||
options.Port,
|
||
options.User,
|
||
options.Password,
|
||
options.DBName,
|
||
options.SSLMode,
|
||
options.TimeZone,
|
||
)
|
||
|
||
// 配置GORM日志（按传入的日志级别，生产环境可关闭）
|
||
gormConfig := &gorm.Config{
|
||
Logger: logger.Default.LogMode(options.LogLevel),
|
||
}
|
||
|
||
// 建立GORM连接
|
||
pgInstance, pgErr = gorm.Open(postgres.Open(dsn), gormConfig)
|
||
if pgErr != nil {
|
||
pgErr = fmt.Errorf("postgres connect failed: %w", pgErr)
|
||
return
|
||
}
|
||
|
||
// 获取底层*sql.DB，配置连接池（关键：避免连接泄漏，提升性能）
|
||
sqlDB, err := pgInstance.DB()
|
||
if err != nil {
|
||
pgErr = fmt.Errorf("get postgres sql.DB failed: %w", err)
|
||
return
|
||
}
|
||
|
||
// 设置连接池参数
|
||
sqlDB.SetMaxOpenConns(options.MaxOpenConns)
|
||
sqlDB.SetMaxIdleConns(options.MaxIdleConns)
|
||
sqlDB.SetConnMaxLifetime(options.ConnMaxLifetime)
|
||
sqlDB.SetConnMaxIdleTime(options.ConnMaxIdleTime)
|
||
|
||
// 测试连接（ping一下，确保连接有效）
|
||
if err := sqlDB.PingContext(context.Background()); err != nil {
|
||
pgErr = fmt.Errorf("postgres ping failed: %w", err)
|
||
pgInstance = nil // 连接失败，置空实例
|
||
return
|
||
}
|
||
})
|
||
|
||
// 单例执行完成后，返回实例和错误
|
||
return pgInstance, pgErr
|
||
}
|
||
|
||
// GetPostgres 获取全局PostgreSQL单例连接
|
||
// 业务层初始化后，可通过此方法直接获取连接，无需重复调用InitPostgres
|
||
func GetPostgres() (*gorm.DB, error) {
|
||
if pgInstance == nil {
|
||
return nil, fmt.Errorf("postgres not initialized, please call InitPostgres first")
|
||
}
|
||
return pgInstance, nil
|
||
}
|
||
|
||
// ClosePostgres 关闭PostgreSQL连接（应用退出时调用，释放资源）
|
||
// 一般在main函数的defer中调用，如：defer db.ClosePostgres()
|
||
func ClosePostgres() error {
|
||
if pgInstance == nil {
|
||
return fmt.Errorf("postgres not initialized")
|
||
}
|
||
|
||
// 获取底层*sql.DB，执行关闭
|
||
sqlDB, err := pgInstance.DB()
|
||
if err != nil {
|
||
return fmt.Errorf("get postgres sql.DB failed: %w", err)
|
||
}
|
||
|
||
if err := sqlDB.Close(); err != nil {
|
||
return fmt.Errorf("close postgres failed: %w", err)
|
||
}
|
||
|
||
// 关闭后重置单例，避免重复关闭
|
||
pgOnce = sync.Once{}
|
||
pgInstance = nil
|
||
pgErr = nil
|
||
|
||
return nil
|
||
}
|
||
|
||
// PingPostgres 测试PostgreSQL连接是否有效
|
||
// 业务层可定时调用（如健康检查），确保连接未断开
|
||
func PingPostgres() error {
|
||
db, err := GetPostgres()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
sqlDB, err := db.DB()
|
||
if err != nil {
|
||
return fmt.Errorf("get postgres sql.DB failed: %w", err)
|
||
}
|
||
|
||
return sqlDB.PingContext(context.Background())
|
||
} |