pgm/qry_insert.go

// Patial Tech.
// Author, Ankit Patial

package pgm

import (
	"context"
	"fmt"
	"strconv"
	"strings"

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

type (
	InsertClause interface {
		Insert
		Returning(field Field) First
		OnConflict(fields ...Field) Do
		Execute
		Stringer
	}

	Insert interface {
		Set(field Field, val any) InsertClause
		SetMap(fields map[Field]any) InsertClause
	}

	insertQry struct {
		returing   *string
		onConflict *string

		table          string
		conflictAction string

		fields []string
		vals   []string
		args   []any
		debug  bool
	}
)

func (q *insertQry) Set(field Field, val any) InsertClause {
	q.fields = append(q.fields, field.Name())
	q.vals = append(q.vals, "$"+strconv.Itoa(len(q.args)+1))
	q.args = append(q.args, val)
	return q
}

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

func (q *insertQry) Returning(field Field) First {
	col := field.String()
	q.returing = &col
	return q
}

func (q *insertQry) OnConflict(fields ...Field) Do {
	if len(fields) > 0 {
		sb := getSB()
		defer putSB(sb)

		for i, f := range fields {
			if i == 0 {
				sb.WriteString(f.Name())
			} else {
				sb.WriteString(", " + f.Name())
			}
		}
		c := sb.String()
		q.onConflict = &c
	}
	return q
}

func (q *insertQry) DoNothing() Execute {
	q.conflictAction = "DO NOTHING"
	return q
}

func (q *insertQry) DoUpdate(fields ...Field) Execute {
	var sb strings.Builder
	for i, f := range fields {
		col := f.Name()
		if i == 0 {
			fmt.Fprintf(&sb, "%s = EXCLUDED.%s", col, col)
		} else {
			fmt.Fprintf(&sb, ", %s = EXCLUDED.%s", col, col)
		}
	}

	q.conflictAction = "DO UPDATE SET " + sb.String()
	return q
}

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

func (q *insertQry) 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 *insertQry) First(ctx context.Context, dest ...any) error {
	return poolPGX.Load().QueryRow(ctx, q.String(), q.args...).Scan(dest...)
}

func (q *insertQry) FirstTx(ctx context.Context, tx pgx.Tx, dest ...any) error {
	return tx.QueryRow(ctx, q.String(), q.args...).Scan(dest...)
}

// build query string
func (q *insertQry) String() string {
	sb := getSB()
	defer putSB(sb)

	n := 12 + len(q.table) + 10
	for i, c := range q.fields {
		n += len(c) + len(" =$,"+strconv.Itoa(i))
	}

	sb.Grow(n)
	sb.WriteString("INSERT INTO ")
	sb.WriteString(q.table)
	sb.WriteString("(")
	sb.WriteString(strings.Join(q.fields, ", "))
	sb.WriteString(") VALUES(")
	sb.WriteString(strings.Join(q.vals, ", "))
	sb.WriteString(")")

	if q.onConflict != nil {
		sb.WriteString(" ON CONFLICT (" + *q.onConflict + ") " + q.conflictAction)
	}

	if q.returing != nil {
		sb.WriteString(" RETURNING " + *q.returing)
	}

	return sb.String()
}