pgm/pool.go

// Patial Tech.
// Author, Ankit Patial

package pgm

import (
	"context"
	"errors"
	"log/slog"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	poolPGX           atomic.Pointer[pgxpool.Pool]
	poolStringBuilder = sync.Pool{
		New: func() any {
			return new(strings.Builder)
		},
	}

	ErrInitTX   = errors.New("failed to init db.tx")
	ErrCommitTX = errors.New("failed to commit db.tx")
	ErrNoRows   = errors.New("no data found")
)

type Config struct {
	MaxConns        int32
	MinConns        int32
	MaxConnLifetime time.Duration
	MaxConnIdleTime time.Duration
}

func Init(connString string, conf *Config) {
	cfg, err := pgxpool.ParseConfig(connString)
	if err != nil {
		panic(err)
	}

	if conf != nil {
		if conf.MaxConns > 0 {
			cfg.MaxConns = conf.MaxConns // 100
		}

		if conf.MinConns > 0 {
			cfg.MinConns = conf.MaxConns // 5
		}

		if conf.MaxConnLifetime > 0 {
			cfg.MaxConnLifetime = conf.MaxConnLifetime // time.Minute * 10
		}

		if conf.MaxConnIdleTime > 0 {
			cfg.MaxConnIdleTime = conf.MaxConnIdleTime // time.Minute * 5
		}
	}

	p, err := pgxpool.NewWithConfig(context.Background(), cfg)
	if err != nil {
		panic(err)
	}

	if err = p.Ping(context.Background()); err != nil {
		panic(err)
	}

	poolPGX.Store(p)
}

func GetPool() *pgxpool.Pool {
	return poolPGX.Load()
}

// get string builder from pool
func getSB() *strings.Builder {
	return poolStringBuilder.Get().(*strings.Builder)
}

// put string builder back to pool
func putSB(sb *strings.Builder) {
	sb.Reset()
	poolStringBuilder.Put(sb)
}

func BeginTx(ctx context.Context) (pgx.Tx, error) {
	tx, err := poolPGX.Load().Begin(ctx)
	if err != nil {
		slog.Error(err.Error())
		return nil, errors.New("failed to open db tx")
	}

	return tx, err
}