From f63c795dcb7b325d188b87d910f46608aa7f8531 Mon Sep 17 00:00:00 2001 From: lan Date: Fri, 15 May 2026 14:44:40 +0800 Subject: [PATCH] feat(db): implement read-write splitting for postgres Add support for database read replicas using the GORM dbresolver plugin. This allows for scaling read operations by distributing them across one or multiple replica nodes. - Update `DatabaseConfig` to support single and multiple replica configurations. - Add configuration options for replica connection pooling and selection policy. - Integrate `dbresolver` in the database initialization process for PostgreSQL. - Add helper methods to aggregate replica configurations. --- configs/config.yaml | 22 +++++++++++++++++++ internal/config/config.go | 1 + internal/config/database.go | 42 ++++++++++++++++++++++++++++--------- internal/model/init.go | 35 +++++++++++++++++++++++++++++++ 4 files changed, 90 insertions(+), 10 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 6115dce..a49c9aa 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -21,6 +21,28 @@ database: password: postgres dbname: with_you sslmode: disable + # 读写分离配置（仅 postgres 模式生效，不配置则所有请求走主库） + # 单副本配置（支持环境变量）: + # APP_DATABASE_REPLICA_HOST, APP_DATABASE_REPLICA_PORT, APP_DATABASE_REPLICA_USER, + # APP_DATABASE_REPLICA_PASSWORD, APP_DATABASE_REPLICA_DBNAME, APP_DATABASE_REPLICA_SSLMODE + # replica: + # host: replica1.internal + # port: 5432 + # user: readonly + # password: readonly_pass + # dbname: with_you + # sslmode: disable + # 多副本配置（仅 YAML，环境变量不支持切片结构体）: + # replicas: + # - host: replica1.internal + # port: 5432 + # user: readonly + # password: readonly_pass + # dbname: with_you + # sslmode: disable + # replica_policy: random # random + # replica_max_idle_conns: 10 + # replica_max_open_conns: 100 max_idle_conns: 10 max_open_conns: 100 log_level: warn diff --git a/internal/config/config.go b/internal/config/config.go index 462b836..d0088a0 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -59,6 +59,7 @@ func Load(configPath string) (*Config, error) { viper.SetDefault("database.slow_threshold_ms", 200) viper.SetDefault("database.ignore_record_not_found", true) viper.SetDefault("database.parameterized_queries", true) + viper.SetDefault("database.replica_policy", "random") viper.SetDefault("redis.type", "miniredis") viper.SetDefault("redis.redis.host", "localhost") viper.SetDefault("redis.redis.port", 6379) diff --git a/internal/config/database.go b/internal/config/database.go index 2dcc14a..08e606e 100644 --- a/internal/config/database.go +++ b/internal/config/database.go @@ -4,15 +4,22 @@ import "fmt" // DatabaseConfig 数据库配置 type DatabaseConfig struct { - Type string `mapstructure:"type"` - SQLite SQLiteConfig `mapstructure:"sqlite"` - Postgres PostgresConfig `mapstructure:"postgres"` - MaxIdleConns int `mapstructure:"max_idle_conns"` - MaxOpenConns int `mapstructure:"max_open_conns"` - LogLevel string `mapstructure:"log_level"` - SlowThresholdMs int `mapstructure:"slow_threshold_ms"` - IgnoreRecordNotFound bool `mapstructure:"ignore_record_not_found"` - ParameterizedQueries bool `mapstructure:"parameterized_queries"` + Type string `mapstructure:"type"` + SQLite SQLiteConfig `mapstructure:"sqlite"` + Postgres PostgresConfig `mapstructure:"postgres"` + MaxIdleConns int `mapstructure:"max_idle_conns"` + MaxOpenConns int `mapstructure:"max_open_conns"` + LogLevel string `mapstructure:"log_level"` + SlowThresholdMs int `mapstructure:"slow_threshold_ms"` + IgnoreRecordNotFound bool `mapstructure:"ignore_record_not_found"` + ParameterizedQueries bool `mapstructure:"parameterized_queries"` + // 读写分离：单副本配置（支持环境变量，如 APP_DATABASE_REPLICA_HOST） + Replica PostgresConfig `mapstructure:"replica"` + // 读写分离：多副本配置（仅 YAML 支持，环境变量无法覆盖切片结构体） + Replicas []PostgresConfig `mapstructure:"replicas"` + ReplicaPolicy string `mapstructure:"replica_policy"` + ReplicaMaxIdle int `mapstructure:"replica_max_idle_conns"` + ReplicaMaxOpen int `mapstructure:"replica_max_open_conns"` } // SQLiteConfig SQLite 配置 @@ -20,7 +27,7 @@ type SQLiteConfig struct { Path string `mapstructure:"path"` } -// PostgresConfig PostgreSQL 配置 +// PostgresConfig PostgreSQL 配置（主库和副本共用） type PostgresConfig struct { Host string `mapstructure:"host"` Port int `mapstructure:"port"` @@ -37,3 +44,18 @@ func (d PostgresConfig) DSN() string { d.Host, d.Port, d.User, d.Password, d.DBName, d.SSLMode, ) } + +// HasReplica 返回是否配置了读副本（单副本或多副本） +func (d DatabaseConfig) HasReplica() bool { + return d.Replica.Host != "" || len(d.Replicas) > 0 +} + +// AllReplicas 返回所有读副本配置（合并单副本 + 多副本） +func (d DatabaseConfig) AllReplicas() []PostgresConfig { + var result []PostgresConfig + if d.Replica.Host != "" { + result = append(result, d.Replica) + } + result = append(result, d.Replicas...) + return result +} \ No newline at end of file diff --git a/internal/model/init.go b/internal/model/init.go index a0981db..712e0a3 100644 --- a/internal/model/init.go +++ b/internal/model/init.go @@ -11,6 +11,7 @@ import ( "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" + "gorm.io/plugin/dbresolver" "with_you/internal/config" ) @@ -63,6 +64,40 @@ func NewDB(cfg *config.DatabaseConfig) (*gorm.DB, error) { sqlDB.SetMaxOpenConns(cfg.MaxOpenConns) } + // 配置读写分离（仅 PostgreSQL 且配置了读副本时生效） + if (cfg.Type == "postgres" || cfg.Type == "postgresql") && cfg.HasReplica() { + allReplicas := cfg.AllReplicas() + replicaDialectors := make([]gorm.Dialector, len(allReplicas)) + for i, replica := range allReplicas { + replicaDialectors[i] = postgres.Open(replica.DSN()) + } + + replicaMaxIdle := cfg.ReplicaMaxIdle + if replicaMaxIdle == 0 { + replicaMaxIdle = cfg.MaxIdleConns + } + replicaMaxOpen := cfg.ReplicaMaxOpen + if replicaMaxOpen == 0 { + replicaMaxOpen = cfg.MaxOpenConns + } + + resolver := dbresolver.Register(dbresolver.Config{ + Replicas: replicaDialectors, + }). + SetMaxIdleConns(replicaMaxIdle). + SetMaxOpenConns(replicaMaxOpen). + SetConnMaxLifetime(time.Hour) + + if err := db.Use(resolver); err != nil { + return nil, fmt.Errorf("failed to register dbresolver plugin: %w", err) + } + + zap.L().Info("Database read/write splitting configured", + zap.String("type", cfg.Type), + zap.Int("replica_count", len(allReplicas)), + ) + } + // 自动迁移 if err := autoMigrate(db); err != nil { return nil, fmt.Errorf("failed to auto migrate: %w", err)