numerus/pkg/db.go

package pkg

import (
	"context"
	"log"

	"github.com/jackc/pgio"
	"github.com/jackc/pgtype"
	"github.com/jackc/pgx/v4"
	"github.com/jackc/pgx/v4/pgxpool"
)

type Db struct {
	*pgxpool.Pool
}

type newInvoiceProductArray []*invoiceProductForm

type discountRate struct {
	pgtype.Numeric
}

func (form *invoiceProductForm) EncodeBinary(ci *pgtype.ConnInfo, buf []byte) (newBuff []byte, err error) {
	productId := &pgtype.Int4{}
	if err = productId.Set(form.ProductId.Val); err != nil {
		return
	}
	name := &pgtype.Text{String: form.Name.Val, Status: pgtype.Present}
	description := &pgtype.Text{String: form.Description.Val, Status: pgtype.Present}
	price := &pgtype.Text{String: form.Price.Val, Status: pgtype.Present}
	quantity := &pgtype.Int4{}
	if err = quantity.Set(form.Quantity.Val); err != nil {
		return
	}
	discount := &discountRate{}
	if err = discount.Set(form.Discount.Val); err != nil {
		return
	}
	tax := &pgtype.Int4Array{}
	if err = tax.Set(form.Tax.Selected); err != nil {
		return
	}
	ct := pgtype.CompositeFields{productId, name, description, price, quantity, discount, tax}
	return ct.EncodeBinary(ci, buf)
}

func (src newInvoiceProductArray) EncodeBinary(ci *pgtype.ConnInfo, buf []byte) ([]byte, error) {
	arrayHeader := pgtype.ArrayHeader{
		Dimensions: []pgtype.ArrayDimension{{Length: int32(len(src)), LowerBound: 1}},
	}

	/*typeName := "numerus.new_invoice_product"
	if dt, ok := ci.DataTypeForName(typeName); ok {
		arrayHeader.ElementOID = int32(dt.OID)
	} else {
		return nil, fmt.Errorf("unable to find oid for type name %v", typeName)
	}*/
	arrayHeader.ElementOID = 69053
	buf = arrayHeader.EncodeBinary(ci, buf)
	for i := range src {
		sp := len(buf)
		buf = pgio.AppendInt32(buf, -1)

		elemBuf, err := src[i].EncodeBinary(ci, buf)
		if err != nil {
			return nil, err
		}
		if elemBuf != nil {
			buf = elemBuf
			pgio.SetInt32(buf[sp:], int32(len(buf[sp:])-4))
		}
	}

	return buf, nil
}

func NewDatabase(ctx context.Context, connString string) (*Db, error) {
	config, err := pgxpool.ParseConfig(connString)
	if err != nil {
		log.Fatal(err)
	}

	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
		_, err := conn.Exec(context.Background(), "SET search_path TO numerus, public")
		conn.ConnInfo().RegisterDataType(pgtype.DataType{Value: &discountRate{}, Name: "discount_rate", OID: 69002})
		return err
	}

	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
		cookie := ""
		if value, ok := ctx.Value(ContextCookieKey).(string); ok {
			cookie = value
		}
		if _, err := conn.Exec(ctx, "select set_cookie($1)", cookie); err != nil {
			log.Printf("ERROR - Failed to set role: %v", err)
			return false
		}
		return true
	}

	config.AfterRelease = func(conn *pgx.Conn) bool {
		if _, err := conn.Exec(context.Background(), "RESET ROLE"); err != nil {
			log.Printf("ERROR - Failed to reset role: %v", err)
			return false
		}
		return true
	}

	pool, err := pgxpool.ConnectConfig(ctx, config)
	if err != nil {
		return nil, err
	}
	return &Db{pool}, nil
}

func notFoundErrorOrPanic(err error) bool {
	if err == pgx.ErrNoRows {
		return true
	}
	if err != nil {
		panic(err)
	}
	return false
}

func (db *Db) Acquire(ctx context.Context) (*Conn, error) {
	conn, err := db.Pool.Acquire(ctx)
	if err != nil {
		return nil, err
	}
	return &Conn{conn}, nil
}

func (db *Db) MustAcquire(ctx context.Context) *Conn {
	conn, err := db.Acquire(ctx)
	if err != nil {
		panic(err)
	}
	return conn
}

type Conn struct {
	*pgxpool.Conn
}

func (c *Conn) MustBegin(ctx context.Context) *Tx {
	tx, err := c.Begin(ctx)
	if err != nil {
		panic(err)
	}
	return &Tx{tx}
}

func (c *Conn) MustGetText(ctx context.Context, def string, sql string, args ...interface{}) string {
	var result string
	if notFoundErrorOrPanic(c.Conn.QueryRow(ctx, sql, args...).Scan(&result)) {
		return def
	}

	return result
}

func (c *Conn) MustGetBool(ctx context.Context, sql string, args ...interface{}) bool {
	var result bool
	if err := c.Conn.QueryRow(ctx, sql, args...).Scan(&result); err != nil {
		panic(err)
	}
	return result
}

func (c *Conn) MustExec(ctx context.Context, sql string, args ...interface{}) {
	if _, err := c.Conn.Exec(ctx, sql, args...); err != nil {
		panic(err)
	}
}

func (c *Conn) MustQuery(ctx context.Context, sql string, args ...interface{}) pgx.Rows {
	rows, err := c.Conn.Query(ctx, sql, args...)
	if err != nil {
		panic(err)
	}
	return rows
}

type Tx struct {
	pgx.Tx
}

func (tx *Tx) MustCommit(ctx context.Context) {
	if err := tx.Commit(ctx); err != nil {
		panic(err)
	}
}

func (tx *Tx) MustRollback(ctx context.Context) {
	if err := tx.Rollback(ctx); err != nil {
		panic(err)
	}
}

func (tx *Tx) MustGetInteger(ctx context.Context, sql string, args ...interface{}) int {
	var result int
	if err := tx.QueryRow(ctx, sql, args...).Scan(&result); err != nil {
		panic(err)
	}
	return result
}

func (tx *Tx) MustGetIntegerOrDefault(ctx context.Context, def int, sql string, args ...interface{}) int {
	var result int
	if notFoundErrorOrPanic(tx.QueryRow(ctx, sql, args...).Scan(&result)) {
		return def
	}
	return result
}

func (tx *Tx) MustCopyFrom(ctx context.Context, tableName string, columns []string, rows [][]interface{}) int64 {
	copied, err := tx.CopyFrom(ctx, pgx.Identifier{tableName}, columns, pgx.CopyFromRows(rows))
	if err != nil {
		panic(err)
	}
	return copied
}