numerus/pkg/db.go

package pkg

import (
	"context"
	"log"
	"net/http"

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

type Db struct {
	pool *pgxpool.Pool
}

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")
		return err
	}

	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
		if user, ok := ctx.Value(ContextUserKey).(*AppUser); ok {
			batch := &pgx.Batch{}
			batch.Queue("select set_config('request.user', $1, false)", user.Email)
			batch.Queue("select set_config('role', $1, false)", user.Role)
			br := conn.SendBatch(ctx, batch)
			defer br.Close()
			for i := 0; i < batch.Len(); i++ {
				if _, err := br.Exec(); 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 (db *Db) Close() {
	db.pool.Close()
}

func (db *Db) Text(r *http.Request, def string, sql string, args ...interface{}) string {
	var result string
	if err := db.pool.QueryRow(r.Context(), sql, args...).Scan(&result); err != nil {
		if err == pgx.ErrNoRows {
			return def
		}
		panic(err)
	}

	return result
}

func (db *Db) Exec(r *http.Request, sql string, args ...interface{}) {
	if _, err := db.pool.Exec(r.Context(), sql, args...); err != nil {
		panic(err)
	}
}