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)
		},
	}

	ErrConnStringMissing = errors.New("connection string is empty")
	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 {
	ConnString      string
	MaxConns        int32
	MinConns        int32
	MaxConnLifetime time.Duration
	MaxConnIdleTime time.Duration
}

// InitPool will create new pgxpool.Pool and will keep it for its working
func InitPool(conf Config) {
	if conf.ConnString == "" {
		panic(ErrConnStringMissing)
	}

	cfg, err := pgxpool.ParseConfig(conf.ConnString)
	if err != nil {
		panic(err)
	}

	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)
}

// 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)
}

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

// BeginTx begins a pgx poll transaction
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
}

// IsNotFound error check
func IsNotFound(err error) bool {
	return errors.Is(err, pgx.ErrNoRows)
}