151 lines
2.9 KiB
Go
151 lines
2.9 KiB
Go
// Patial Tech.
|
|
// Author, Ankit Patial
|
|
|
|
package pgm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
type insertQry struct {
|
|
returing *string
|
|
onConflict *string
|
|
|
|
table string
|
|
conflictAction string
|
|
|
|
fields []string
|
|
vals []string
|
|
args []any
|
|
debug bool
|
|
}
|
|
|
|
func (t *Table) Insert() Insert {
|
|
qb := &insertQry{
|
|
table: t.Name,
|
|
fields: make([]string, 0, t.FieldCount),
|
|
vals: make([]string, 0, t.FieldCount),
|
|
args: make([]any, 0, t.FieldCount),
|
|
debug: t.debug,
|
|
}
|
|
return qb
|
|
}
|
|
|
|
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.Name()
|
|
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()
|
|
}
|