93 lines
2.3 KiB
Go
93 lines
2.3 KiB
Go
package database
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
// PostgreSQLOptions defines options for PostgreSQL database.
|
|
type PostgreSQLOptions struct {
|
|
Addr string
|
|
Username string
|
|
Password string
|
|
Database string
|
|
MaxIdleConnections int
|
|
MaxOpenConnections int
|
|
MaxConnectionLifeTime time.Duration
|
|
// +optional
|
|
Logger logger.Interface
|
|
}
|
|
|
|
// DSN return DSN from PostgreSQLOptions.
|
|
func (o *PostgreSQLOptions) DSN() string {
|
|
splited := strings.Split(o.Addr, ":")
|
|
host, port := splited[0], "5432"
|
|
if len(splited) > 1 {
|
|
port = splited[1]
|
|
}
|
|
|
|
return fmt.Sprintf(`user=%s password=%s host=%s port=%s dbname=%s sslmode=disable TimeZone=Asia/Shanghai`,
|
|
o.Username,
|
|
o.Password,
|
|
host,
|
|
port,
|
|
o.Database,
|
|
)
|
|
}
|
|
|
|
// NewPostgreSQL create a new gorm db instance with the given options.
|
|
func NewPostgreSQL(opts *PostgreSQLOptions) (*gorm.DB, error) {
|
|
// Set default values to ensure all fields in opts are available.
|
|
setPostgreSQLDefaults(opts)
|
|
|
|
db, err := gorm.Open(postgres.Open(opts.DSN()), &gorm.Config{
|
|
// PrepareStmt executes the given query in cached statement.
|
|
// This can improve performance.
|
|
PrepareStmt: true,
|
|
Logger: opts.Logger,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// SetMaxOpenConns sets the maximum number of open connections to the database.
|
|
sqlDB.SetMaxOpenConns(opts.MaxOpenConnections)
|
|
|
|
// SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
|
|
sqlDB.SetConnMaxLifetime(opts.MaxConnectionLifeTime)
|
|
|
|
// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
|
|
sqlDB.SetMaxIdleConns(opts.MaxIdleConnections)
|
|
|
|
return db, nil
|
|
}
|
|
|
|
// setPostgreSQLDefaults set available default values for some fields.
|
|
func setPostgreSQLDefaults(opts *PostgreSQLOptions) {
|
|
if opts.Addr == "" {
|
|
opts.Addr = "127.0.0.1:5432"
|
|
}
|
|
if opts.MaxIdleConnections == 0 {
|
|
opts.MaxIdleConnections = 100
|
|
}
|
|
if opts.MaxOpenConnections == 0 {
|
|
opts.MaxOpenConnections = 100
|
|
}
|
|
if opts.MaxConnectionLifeTime == 0 {
|
|
opts.MaxConnectionLifeTime = time.Duration(10) * time.Second
|
|
}
|
|
if opts.Logger == nil {
|
|
opts.Logger = logger.Default.LogMode(logger.Info)
|
|
}
|
|
}
|