init

2025-12-18 15:11:02 +08:00
commit 96858c70e6
9 changed files with 1223 additions and 0 deletions
--- a/database.go
+++ b/database.go
@@ -0,0 +1,170 @@
+package reconnect
+
+import (
+	"context"
+	"database/sql"
+	"sync"
+	"time"
+
+	"gorm.io/gorm"
+)
+
+// DatabaseConfig 数据库配置
+type DatabaseConfig struct {
+	// Dialector gorm dialector（如 postgres.Open(dsn), mysql.Open(dsn)）
+	Dialector gorm.Dialector
+	// GormConfig gorm配置
+	GormConfig *gorm.Config
+	// MaxIdleConns 最大空闲连接数
+	MaxIdleConns int
+	// MaxOpenConns 最大打开连接数
+	MaxOpenConns int
+	// ConnMaxLifetime 连接最大生命周期
+	ConnMaxLifetime time.Duration
+	// ReconnectConfig 重连配置
+	ReconnectConfig Config
+}
+
+// DatabaseClient 带重连功能的数据库客户端
+type DatabaseClient struct {
+	config  DatabaseConfig
+	db      *gorm.DB
+	sqlDB   *sql.DB
+	manager *ConnectionManager
+	mu      sync.RWMutex
+}
+
+// databaseConnector 数据库连接器
+type databaseConnector struct {
+	client *DatabaseClient
+}
+
+func (d *databaseConnector) Connect(ctx context.Context) error {
+	d.client.mu.Lock()
+	defer d.client.mu.Unlock()
+
+	gormConfig := d.client.config.GormConfig
+	if gormConfig == nil {
+		gormConfig = &gorm.Config{}
+	}
+
+	db, err := gorm.Open(d.client.config.Dialector, gormConfig)
+	if err != nil {
+		return err
+	}
+
+	sqlDB, err := db.DB()
+	if err != nil {
+		return err
+	}
+
+	// 配置连接池
+	if d.client.config.MaxIdleConns > 0 {
+		sqlDB.SetMaxIdleConns(d.client.config.MaxIdleConns)
+	}
+	if d.client.config.MaxOpenConns > 0 {
+		sqlDB.SetMaxOpenConns(d.client.config.MaxOpenConns)
+	}
+	if d.client.config.ConnMaxLifetime > 0 {
+		sqlDB.SetConnMaxLifetime(d.client.config.ConnMaxLifetime)
+	}
+
+	// 验证连接
+	if err := sqlDB.PingContext(ctx); err != nil {
+		sqlDB.Close()
+		return err
+	}
+
+	d.client.db = db
+	d.client.sqlDB = sqlDB
+	return nil
+}
+
+func (d *databaseConnector) Close() error {
+	d.client.mu.Lock()
+	defer d.client.mu.Unlock()
+
+	if d.client.sqlDB != nil {
+		err := d.client.sqlDB.Close()
+		d.client.db = nil
+		d.client.sqlDB = nil
+		return err
+	}
+	return nil
+}
+
+// databaseHealthChecker 数据库健康检查器
+type databaseHealthChecker struct {
+	client *DatabaseClient
+}
+
+func (d *databaseHealthChecker) HealthCheck(ctx context.Context) error {
+	d.client.mu.RLock()
+	sqlDB := d.client.sqlDB
+	d.client.mu.RUnlock()
+
+	if sqlDB == nil {
+		return context.Canceled
+	}
+
+	return sqlDB.PingContext(ctx)
+}
+
+// NewDatabaseClient 创建带重连功能的数据库客户端
+func NewDatabaseClient(cfg DatabaseConfig) (*DatabaseClient, error) {
+	client := &DatabaseClient{
+		config: cfg,
+	}
+
+	connector := &databaseConnector{client: client}
+	checker := &databaseHealthChecker{client: client}
+
+	client.manager = NewConnectionManager(connector, checker, cfg.ReconnectConfig)
+
+	// 首次连接
+	ctx := context.Background()
+	if err := client.manager.ConnectWithRetry(ctx); err != nil {
+		return nil, err
+	}
+
+	return client, nil
+}
+
+// GetDB 获取gorm.DB实例
+func (c *DatabaseClient) GetDB() *gorm.DB {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.db
+}
+
+// GetSqlDB 获取sql.DB实例
+func (c *DatabaseClient) GetSqlDB() *sql.DB {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.sqlDB
+}
+
+// State 获取连接状态
+func (c *DatabaseClient) State() ConnectionState {
+	return c.manager.State()
+}
+
+// Close 关闭客户端
+func (c *DatabaseClient) Close() error {
+	return c.manager.Close()
+}
+
+// TriggerReconnect 手动触发重连
+func (c *DatabaseClient) TriggerReconnect() {
+	c.manager.TriggerReconnect()
+}
+
+// Ping 执行Ping命令
+func (c *DatabaseClient) Ping(ctx context.Context) error {
+	sqlDB := c.GetSqlDB()
+	if sqlDB == nil {
+		return context.Canceled
+	}
+	return sqlDB.PingContext(ctx)
+}
+