m

code/base-project/base-go/pkg/config/database.go

@@ -0,0 +1,189 @@
+package config
+
+import (
+	"base-framework/pkg/types"
+	"database/sql"
+	"fmt"
+	"log"
+	"sync"
+	"time"
+
+	"github.com/fsnotify/fsnotify"
+	"github.com/spf13/viper"
+
+	_ "github.com/lib/pq" // PostgreSQL 驱动，根据实际数据库替换
+)
+
+// 默认连接池参数
+const (
+	defaultMaxOpenConns    = 10
+	defaultMaxIdleConns    = 5
+	defaultConnMaxLifetime = time.Hour
+)
+
+var (
+	dbConfigs   map[string]types.DBConfig  // 配置副本
+	dataSources = make(map[string]*sql.DB) // 连接池
+	configsDSN  = make(map[string]string)  // key -> dsn
+	mu          sync.RWMutex
+	v           *viper.Viper
+)
+
+// InitDBConfig 初始化并加载 db.yaml 配置，同时启动监听
+func InitDBConfig(configPath string) error {
+	v = viper.New()
+	v.SetConfigFile(configPath)
+	v.SetConfigType("yaml")
+
+	if err := v.ReadInConfig(); err != nil {
+		return fmt.Errorf("读取数据库配置失败: %w", err)
+	}
+
+	if err := unmarshalConfigs(); err != nil {
+		return err
+	}
+
+	// 首次加载后，初始化数据源
+	if err := ReloadDataSources(GetDBConfigs()); err != nil {
+		log.Printf("首次加载数据源失败: %v\n", err)
+		return err
+	}
+
+	// 监听配置文件变化
+	v.WatchConfig()
+	v.OnConfigChange(func(e fsnotify.Event) {
+		log.Printf("数据库配置文件发生变化: %s\n", e.Name)
+		if err := unmarshalConfigs(); err != nil {
+			log.Printf("重新加载数据库配置失败: %v\n", err)
+		} else {
+			if err := ReloadDataSources(GetDBConfigs()); err != nil {
+				log.Printf("更新数据源失败: %v\n", err)
+			} else {
+				log.Println("数据库连接池已更新")
+			}
+		}
+	})
+
+	return nil
+}
+
+// unmarshalConfigs 解析配置到全局变量，线程安全
+func unmarshalConfigs() error {
+	mu.Lock()
+	defer mu.Unlock()
+
+	temp := make(map[string]types.DBConfig)
+	if err := v.Unmarshal(&temp); err != nil {
+		return fmt.Errorf("解析数据库配置失败: %w", err)
+	}
+	dbConfigs = temp
+	return nil
+}
+
+// GetDBConfigs 并发安全地返回当前数据库配置副本
+func GetDBConfigs() map[string]types.DBConfig {
+	mu.RLock()
+	defer mu.RUnlock()
+
+	res := make(map[string]types.DBConfig, len(dbConfigs))
+	for k, v := range dbConfigs {
+		res[k] = v
+	}
+	return res
+}
+
+// buildDSN 构建 PostgreSQL DSN
+func buildDSN(cfg types.DBConfig) string {
+	return fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+		cfg.Host, cfg.Port, cfg.Username, cfg.Password, cfg.Dbname)
+}
+
+// ReloadDataSources 支持增量更新数据源，并验证数据库连接
+func ReloadDataSources(newConfigs map[string]types.DBConfig) error {
+	mu.Lock()
+	defer mu.Unlock()
+
+	retain := make(map[string]bool)
+
+	for key, cfg := range newConfigs {
+		newDSN := buildDSN(cfg)
+		oldDSN, exists := configsDSN[key]
+
+		// 新增数据源或配置变更
+		if !exists || oldDSN != newDSN {
+			// 关闭旧连接（如果存在）
+			if exists {
+				if oldDB := dataSources[key]; oldDB != nil {
+					_ = oldDB.Close()
+				}
+			}
+
+			db, err := sql.Open("postgres", newDSN)
+			if err != nil {
+				log.Printf("[db] 数据源 %s Open 失败: %v", key, err)
+				continue
+			}
+
+			// 尝试实际连接数据库
+			if err := db.Ping(); err != nil {
+				log.Printf("[db] 数据源 %s 连接失败: %v", key, err)
+				_ = db.Close()
+				continue
+			}
+
+			applyDBConfig(db, cfg)
+			dataSources[key] = db
+			configsDSN[key] = newDSN
+
+			if !exists {
+				log.Printf("[db] 新增数据源 %s 成功", key)
+			} else {
+				log.Printf("[db] 更新数据源 %s 成功", key)
+			}
+		}
+
+		// 配置没变，不用操作
+		retain[key] = true
+	}
+
+	// 关闭并移除不再需要的数据源
+	for key, db := range dataSources {
+		if !retain[key] {
+			_ = db.Close()
+			delete(dataSources, key)
+			delete(configsDSN, key)
+			log.Printf("[db] 移除数据源 %s", key)
+		}
+	}
+
+	return nil
+}
+
+// applyDBConfig 设置 sql.DB 连接池参数，并应用默认值
+func applyDBConfig(db *sql.DB, cfg types.DBConfig) {
+	maxOpen := cfg.MaxOpenConns
+	if maxOpen <= 0 {
+		maxOpen = defaultMaxOpenConns
+	}
+	db.SetMaxOpenConns(maxOpen)
+
+	maxIdle := cfg.MaxIdleConns
+	if maxIdle <= 0 {
+		maxIdle = defaultMaxIdleConns
+	}
+	db.SetMaxIdleConns(maxIdle)
+
+	timeout := cfg.ConnMaxLifetime
+	if timeout <= 0 {
+		timeout = defaultConnMaxLifetime
+	}
+	db.SetConnMaxLifetime(timeout)
+}
+
+// GetDB 获取数据库连接
+func GetDB(key string) (*sql.DB, bool) {
+	mu.RLock()
+	defer mu.RUnlock()
+	db, ok := dataSources[key]
+	return db, ok
+}