1

internal/pkg/database/postgresql.go

@@ -0,0 +1,92 @@
+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
+	}
+}