add

trading_assistant_api/common/db/postgres.go

@@ -0,0 +1,136 @@
+// 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())
+}