80 lines
1.8 KiB
Go
80 lines
1.8 KiB
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)
|
|
}
|
|
}
|