pgm/qry.go

// Patial Tech.
// Author, Ankit Patial

package pgm

import (
	"context"
	"strconv"
	"strings"
	"sync"

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

type (
	Clause interface {
		Insert() InsertClause
		Select(fields ...Field) SelectClause
		Update() UpdateClause
		Delete() DeleteCluase
	}

	Execute interface {
		Exec(ctx context.Context) error
		ExecTx(ctx context.Context, tx pgx.Tx) error
		Stringer
	}

	Stringer interface {
		String() string
	}

	Conditioner interface {
		Condition(args *[]any, idx int) string
	}
)

var sbPool = sync.Pool{
	New: func() any {
		return new(strings.Builder)
	},
}

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

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

func And(cond ...Conditioner) Conditioner {
	return &CondGroup{op: " AND ", cond: cond}
}

func Or(cond ...Conditioner) Conditioner {
	return &CondGroup{op: " OR ", cond: cond}
}

func (cv *Cond) Condition(args *[]any, argIdx int) string {
	// 1. condition with subquery
	if cv.action == CondActionSubQuery {
		qStr, newArgs := cv.Val.(SelectClause).raw(*args)
		*args = newArgs
		return cv.Field + strings.Replace(cv.op, "$", qStr, 1)
	}

	// 2. normal condition
	var op string
	if cv.Val != nil {
		*args = append(*args, cv.Val)
		if strings.HasSuffix(cv.op, "$") {
			op = cv.op + strconv.Itoa(argIdx+1)
		} else {
			op = strings.Replace(cv.op, "$", "$"+strconv.Itoa(argIdx+1), 1)
		}
	} else {
		op = cv.op
	}

	if cv.action == CondActionNeedToClose {
		return cv.Field + op + ")"
	}
	return cv.Field + op
}

func (c *CondGroup) Condition(args *[]any, argIdx int) string {
	sb := getSB()
	defer putSB(sb)

	sb.WriteString("(")
	for i, cond := range c.cond {
		if i == 0 {
			sb.WriteString(cond.Condition(args, argIdx+i))
		} else {
			sb.WriteString(c.op)
			sb.WriteString(cond.Condition(args, argIdx+i))
		}
	}
	sb.WriteString(")")
	return sb.String()
}