golang获取postgres或clickhouse连接

package edgedb

import (
	"database/sql"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"gorm.io/driver/clickhouse"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
)

var (
	pgConn *gorm.DB
	ckConn *gorm.DB
)

type Config struct {
	// 遥感数据存储，(默认: postgresql)
	// 支持 postgresql/clickhouse，可多选通过 "," 连接多个选择
	DeviceTelemetryRecordBackend string `yaml:"" env:"DEVICE_TELEMETRY_RECORD_BACKEND" envDefault:"postgresql"`

	Postgresql struct {
		Host       string `yaml:"host" env:"POSTGRESQL_SERVER_HOST" envDefault:"postgresql"`
		Port       int    `yaml:"port" env:"POSTGRESQL_SERVER_PORT" envDefault:"5432"`
		EdgedbUser string `yaml:"user" env:"POSTGRESQL_SERVER_EDGEDB_USER" envDefault:"postgres"`
		EdgedbPass string `yaml:"pass" env:"POSTGRESQL_SERVER_EDGEDB_PASS" envDefault:"password"`
		EdgedbName string `yaml:"name" env:"POSTGRESQL_SERVER_EDGEDB_NAME" envDefault:"edgedb"`
		PoolSize   int    `yaml:"pool_size" env:"POSTGRESQL_SERVER_POOLSIZE" envDefault:"500"` // 设置连接池的最大连接数
	}

	Clickhouse struct {
		Enabled      bool   `yaml:"host" env:"CLICKHOUSE_SERVER_ENABLED" envDefault:"false"`
		Host         string `yaml:"host" env:"CLICKHOUSE_SERVER_HOST" envDefault:"clickhouse"`
		Port         int    `yaml:"port" env:"CLICKHOUSE_SERVER_PORT" envDefault:"9000"`
		User         string `yaml:"port" env:"CLICKHOUSE_SERVER_PORT" envDefault:"default"`
		Pass         string `yaml:"port" env:"CLICKHOUSE_SERVER_PORT" envDefault:""`
		EdgedbUser   string `yaml:"user" env:"CLICKHOUSE_SERVER_EDGEDB_USER" envDefault:"root"`
		EdgedbPass   string `yaml:"pass" env:"CLICKHOUSE_SERVER_EDGEDB_PASS" envDefault:"[yz^_^edge]"`
		EdgedbName   string `yaml:"name" env:"CLICKHOUSE_SERVER_EDGEDB_NAME" envDefault:"edgedb"`
		ReadTimeout  int    `yaml:"read_timeout" env:"READ_TIMEOUT" envDefault:"10"`
		WriteTimeout int    `yaml:"write_timeout" env:"WRITE_TIMEOUT" envDefault:"10"`
	}
}

// SetupConn 初始化连接，配置连接池
// 注意：不要关闭数据库连接客户端 sql.DB，否则连接池无法生效，如无特殊需求，只需判断有无异常即可。通过 GetXXConn 函数向连接池申请连接
func SetupConn(cfg *Config) (*sql.DB, error) {
	// 加载配置
	if cfg == nil {
		err := errors.New("Config needed")
		return nil, err
	}

	// 初始化 Postgres 连接
	db, err := _setupPostgresConn(cfg)
	if err != nil {
		return nil, err
	}
	defer db.Close()

	// 如系统配置 Clickhouse 作为遥测扩展存储，则额外增加 CK 连接的初始化
	telemetryDatabaseBackends := strings.Split(cfg.DeviceTelemetryRecordBackend, ",")
	sort.Strings(telemetryDatabaseBackends)
	index := sort.SearchStrings(telemetryDatabaseBackends, "clickhouse")
	if index < len(telemetryDatabaseBackends) && telemetryDatabaseBackends[index] == "clickhouse" {
		// 初始化 Clickhouse 连接
		if _, err := _setupClickhouseConn(cfg); err != nil {
			return nil, err
		}
	}

	return db, nil
}

func _setupPostgresConn(cfg *Config) (*sql.DB, error) {
	if cfg == nil {
		if config, err := configs.LoadConfig(); err != nil {
			return nil, err
		} else {
			cfg = config
		}
	}

	dsn := fmt.Sprintf(
		"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable TimeZone=Asia/Shanghai",
		cfg.Postgresql.Host,
		cfg.Postgresql.Port,
		cfg.Postgresql.EdgedbUser,
		cfg.Postgresql.EdgedbPass,
		cfg.Postgresql.EdgedbName,
	)
	conn, err := gorm.Open(postgres.New(postgres.Config{
		DSN:                  dsn,
		PreferSimpleProtocol: true, // disables implicit prepared statement usage
	}), &gorm.Config{})
	if err != nil {
		return nil, err
	}

	pgDB, err := conn.DB()
	if err != nil {
		return nil, err
	}

	// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
	// 设置可空闲的最大连接数，随时等待调用
	pgDB.SetMaxIdleConns(10)

	// SetMaxOpenConns sets the maximum number of open connections to the database.
	// 设置连接池的最大连接数，不配置默认就是不限制 (当前500，依赖 postgres max_connections)
	pgDB.SetMaxOpenConns(cfg.Postgresql.PoolSize)

	// SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
	// 连接的最长存活期，超过这个时间，连接将会重置，不再被复用，不配置默认就是永不过期
	pgDB.SetConnMaxLifetime(time.Second * 600)

	// Global
	pgConn = conn

	return pgDB, nil
}

// 初始化clickhouse的连接，配置连接池
func _setupClickhouseConn(cfg *Config) (*sql.DB, error) {
	if cfg == nil {
		if config, err := configs.LoadConfig(); err != nil {
			return nil, err
		} else {
			cfg = config
		}
	}

	dsn := fmt.Sprintf(
		"tcp://%s:%d?database=%s&username=%s&password=%s&read_timeout=%d&write_timeout=%d",
		cfg.Clickhouse.Host,
		cfg.Clickhouse.Port,
		cfg.Clickhouse.EdgedbName,
		cfg.Clickhouse.EdgedbUser,
		cfg.Clickhouse.EdgedbPass,
		cfg.Clickhouse.ReadTimeout,
		cfg.Clickhouse.WriteTimeout,
	)
	conn, err := gorm.Open(clickhouse.New(clickhouse.Config{
		DSN:                       dsn,
		DisableDatetimePrecision:  true,     // disable datetime64 precision, not supported before clickhouse 20.4
		DontSupportRenameColumn:   true,     // rename column not supported before clickhouse 20.4
		SkipInitializeWithVersion: false,    // smart configure based on used version
		DefaultGranularity:        3,        // 1 granule = 8192 rows
		DefaultCompression:        "LZ4",    // default compression algorithm. LZ4 is lossless
		DefaultIndexType:          "minmax", // index stores extremes of the expression
		DefaultTableEngineOpts:    "ENGINE=MergeTree() ORDER BY tuple()",
	}), &gorm.Config{})
	if err != nil {
		return nil, err
	}
	ckDB, err := conn.DB()
	if err != nil {
		return nil, err
	}

	// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
	// 设置可空闲的最大连接数，随时等待调用
	ckDB.SetMaxIdleConns(1)

	// SetMaxOpenConns sets the maximum number of open connections to the database.
	// 设置连接池的最大连接数，不配置，默认为 0，就是不限制
	ckDB.SetMaxOpenConns(10)

	// SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
	// 连接的最长存活期，超过这个时间，连接将会重置，不再被复用，不配置默认就是永不过期
	ckDB.SetConnMaxLifetime(time.Second * 600)

	ckConn = conn
	return ckDB, nil
}

// GetDBConn 从连接池获取数据库连接
func GetDBConn() (*gorm.DB, error) {
	var pgDB *sql.DB

	if pgConn == nil {
		if db, err := _setupPostgresConn(nil); err != nil {
			return nil, err
		} else {
			pgDB = db
		}
	} else if db, err := pgConn.DB(); err != nil {
		// 尝试获取连接池失败，重新建立连接
		if db, err := _setupPostgresConn(nil); err != nil {
			return nil, err
		} else {
			pgDB = db
		}
	} else {
		pgDB = db
	}

	if err := pgDB.Ping(); err != nil {
		// 尝试 PING 失败，重新建立连接
		_ = pgDB.Close()

		if db, err := _setupPostgresConn(nil); err != nil {
			return nil, err
		} else {
			pgDB = db
		}
	}

	return pgConn, nil
}

// GetDBConn 从连接池获取数据库连接
func GetClickhouseConn() (*gorm.DB, error) {
	var ckDB *sql.DB

	if ckConn == nil {
		if db, err := _setupClickhouseConn(nil); err != nil {
			return nil, err
		} else {
			ckDB = db
		}
	} else if db, err := ckConn.DB(); err != nil {
		if db, err := _setupClickhouseConn(nil); err != nil {
			return nil, err
		} else {
			ckDB = db
		}
	} else {
		ckDB = db
	}

	if err := ckDB.Ping(); err != nil {
		_ = ckDB.Close()

		if db, err := _setupClickhouseConn(nil); err != nil {
			return nil, err
		} else {
			ckDB = db
		}
	}

	return ckConn, nil
}