pgm/qry_update.go

// Patial Tech.
// Author, Ankit Patial

package pgm

import (
	"context"
	"strconv"
	"strings"

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

type updateQry struct {
	table     string
	cols      []string
	condition []Conditioner
	args      []any
	debug     bool
}

func (t *Table) Update() Update {
	qb := &updateQry{
		table: t.Name,
		debug: t.debug,
		cols:  make([]string, 0, t.FieldCount),
		args:  make([]any, 0, t.FieldCount),
	}
	return qb
}

func (q *updateQry) Set(field Field, val any) UpdateClause {
	col := field.Name()
	q.cols = append(q.cols, col+"=$"+strconv.Itoa(len(q.args)+1))
	q.args = append(q.args, val)
	return q
}

func (q *updateQry) SetMap(cols map[Field]any) UpdateClause {
	for k, v := range cols {
		q.Set(k, v)
	}
	return q
}

func (q *updateQry) Where(cond ...Conditioner) WhereOrExec {
	q.condition = append(q.condition, cond...)
	return q
}

func (q *updateQry) Exec(ctx context.Context) error {
	_, err := poolPGX.Load().Exec(ctx, q.String(), q.args...)
	if err != nil {
		return err
	}

	return nil
}

func (q *updateQry) ExecTx(ctx context.Context, tx pgx.Tx) error {
	_, err := tx.Exec(ctx, q.String(), q.args...)
	if err != nil {
		return err
	}

	return nil
}

func (q *updateQry) String() string {
	sb := getSB()
	defer putSB(sb)

	n := 7 + len(q.table) + 5 // "UPDATE q.table SET
	for _, col := range q.cols {
		n += len(col) + 5
	}
	if len(q.condition) > 0 {
		n += 7 + len(q.condition)*5 // WHERE with condition, 5 is just avg small min val
	}
	sb.Grow(n)

	// UPDATE
	sb.WriteString("UPDATE " + q.table + " SET ")
	sb.WriteString(strings.Join(q.cols, ", "))

	// WHERE
	if len(q.condition) > 0 {
		sb.WriteString(" WHERE ")
		var argIdx int
		for i, c := range q.condition {
			argIdx = len(q.args)
			if i > 0 {
				sb.WriteString(" AND ")
			}
			sb.WriteString(c.Condition(&q.args, argIdx))
		}
	}

	return sb.String()
}