camper/pkg/database/db.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
}