127 lines
2.8 KiB
Go
127 lines
2.8 KiB
Go
/*
|
|
* SPDX-FileCopyrightText: 2023 jordi fita mas <jfita@peritasoft.com>
|
|
* SPDX-License-Identifier: AGPL-3.0-only
|
|
*/
|
|
|
|
package database
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log"
|
|
|
|
"github.com/jackc/pgconn"
|
|
"github.com/jackc/pgx/v4"
|
|
"github.com/jackc/pgx/v4/pgxpool"
|
|
)
|
|
|
|
const (
|
|
searchPathQuery = "set search_path to camper, public"
|
|
)
|
|
|
|
func ErrorIsNotFound(err error) bool {
|
|
return errors.Is(err, pgx.ErrNoRows)
|
|
}
|
|
|
|
func New(ctx context.Context, connString string) (*DB, error) {
|
|
// Connect once with an “elevated” user to register types in camper schema
|
|
if err := registerTypes(ctx, connString); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
config, err := pgxpool.ParseConfig(connString)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
|
|
if _, err := conn.Exec(ctx, searchPathQuery); err != nil {
|
|
return err
|
|
}
|
|
return registerConnectionTypes(ctx, conn)
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
type DB struct {
|
|
*pgxpool.Pool
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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) MustExec(ctx context.Context, sql string, args ...interface{}) pgconn.CommandTag {
|
|
tag, err := c.Conn.Exec(ctx, sql, args...)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return tag
|
|
}
|
|
|
|
func (c *Conn) GetText(ctx context.Context, sql string, args ...interface{}) (string, error) {
|
|
var result string
|
|
if err := c.QueryRow(ctx, sql, args...).Scan(&result); err != nil {
|
|
return "", err
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (c *Conn) MustGetText(ctx context.Context, sql string, args ...interface{}) string {
|
|
if result, err := c.GetText(ctx, sql, args...); err == nil {
|
|
return result
|
|
} else {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func (c *Conn) GetBool(ctx context.Context, sql string, args ...interface{}) (bool, error) {
|
|
var result bool
|
|
if err := c.QueryRow(ctx, sql, args...).Scan(&result); err != nil {
|
|
return false, err
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (c *Conn) GetInt(ctx context.Context, sql string, args ...interface{}) (int, error) {
|
|
var result int
|
|
err := c.QueryRow(ctx, sql, args...).Scan(&result)
|
|
return result, err
|
|
}
|
|
|
|
func (c *Conn) GetBytes(ctx context.Context, sql string, args ...interface{}) ([]byte, error) {
|
|
var result []byte
|
|
err := c.QueryRow(ctx, sql, args...).Scan(&result)
|
|
return result, err
|
|
}
|