Merge pull request #28 from interline-io/db-mw

Improve DB and River init

Author: irees

dbutil/db.go

@@ -2,6 +2,7 @@ package dbutil
 
 import (
 	"context"
+	"database/sql"
 	"regexp"
 	"strings"
 	"time"
@@ -11,7 +12,6 @@ import (
 	"github.com/interline-io/log"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jackc/pgx/v5/stdlib"
-	_ "github.com/jackc/pgx/v5/stdlib"
 	"github.com/jmoiron/sqlx"
 	"github.com/jmoiron/sqlx/reflectx"
 )
@@ -25,38 +25,51 @@ func toSnakeCase(str string) string {
 	return strings.ToLower(snake)
 }
 
-func OpenDBPool(ctx context.Context, url string) (*pgxpool.Pool, *sqlx.DB, error) {
-	pool, err := pgxpool.New(ctx, url)
-	if err != nil {
-		return nil, nil, err
-	}
-	db := sqlx.NewDb(stdlib.OpenDBFromPool(pool), "pgx")
+// ConfigureDB sets up common database configuration
+func ConfigureDB(sqlDb *sql.DB) (*sqlx.DB, error) {
+	db := sqlx.NewDb(sqlDb, "pgx")
 	db.SetMaxOpenConns(10)
 	db.SetMaxIdleConns(10)
 	db.SetConnMaxLifetime(time.Hour)
 	if err := db.Ping(); err != nil {
 		log.Error().Err(err).Msgf("could not connect to database")
-		return nil, nil, err
+		return nil, err
 	}
 	db.Mapper = reflectx.NewMapperFunc("db", toSnakeCase)
-	return pool, db.Unsafe(), nil
+	return db.Unsafe(), nil
 }
 
-func OpenDB(url string) (*sqlx.DB, error) {
-	db, err := sqlx.Open("pgx", url)
+func OpenPool(ctx context.Context, url string) (*pgxpool.Pool, error) {
+	pool, err := pgxpool.New(ctx, url)
 	if err != nil {
-		log.Error().Err(err).Msg("could not open database")
+		log.Error().Err(err).Msg("could not open database pool")
 		return nil, err
 	}
-	db.SetMaxOpenConns(10)
-	db.SetMaxIdleConns(10)
-	db.SetConnMaxLifetime(time.Hour)
-	if err := db.Ping(); err != nil {
-		log.Error().Err(err).Msgf("could not connect to database")
+	if err := pool.Ping(ctx); err != nil {
+		log.Error().Err(err).Msg("could not ping database pool")
 		return nil, err
 	}
-	db.Mapper = reflectx.NewMapperFunc("db", toSnakeCase)
-	return db.Unsafe(), nil
+	return pool, nil
+}
+
+func OpenDBPool(ctx context.Context, url string) (*pgxpool.Pool, *sqlx.DB, error) {
+	pool, err := pgxpool.New(ctx, url)
+	if err != nil {
+		return nil, nil, err
+	}
+	db, err := ConfigureDB(stdlib.OpenDBFromPool(pool))
+	if err != nil {
+		return nil, nil, err
+	}
+	return pool, db, nil
+}
+
+func OpenDB(url string) (*sqlx.DB, error) {
+	db, err := sqlx.Open("pgx", url)
+	if err != nil {
+		return nil, err
+	}
+	return ConfigureDB(db.DB)
 }
 
 // Select runs a query and reads results into dest.

jobs/river/river_jobs.go

@@ -43,7 +43,7 @@ func (r riverJobArgs) ToJob() jobs.Job {
 	}
 }
 
-func newRiverJobArgsFrmoJob(job jobs.Job) riverJobArgs {
+func newRiverJobArgsFromJob(job jobs.Job) riverJobArgs {
 	return riverJobArgs{
 		Queue:       job.Queue,
 		JobType:     job.JobType,
@@ -67,28 +67,34 @@ type RiverJobs struct {
 }
 
 func NewRiverJobs(pool *pgxpool.Pool, queuePrefix string) (*RiverJobs, error) {
+	return NewRiverJobsWithMiddleware(pool, queuePrefix)
+}
+
+func NewRiverJobsWithMiddleware(pool *pgxpool.Pool, queuePrefix string, middlewares ...rivertype.Middleware) (*RiverJobs, error) {
 	w := &RiverJobs{
 		pool:        pool,
 		jobMapper:   jobs.NewJobMapper(),
 		queuePrefix: queuePrefix,
 	}
-	return w, w.initClient()
+	return w, w.initClient(middlewares...)
 }
 
 func (w *RiverJobs) RiverClient() *river.Client[pgx.Tx] {
 	return w.riverClient
 }
 
-func (w *RiverJobs) initClient() error {
+func (w *RiverJobs) initClient(middlewares ...rivertype.Middleware) error {
 	var err error
 	defaultQueue := w.queueName("default")
 	w.riverWorkers = river.NewWorkers()
+
 	w.riverClient, err = river.NewClient(riverpgxv5.New(w.pool), &river.Config{
 		Queues:            map[string]river.QueueConfig{defaultQueue: {MaxWorkers: 4}},
 		JobTimeout:        120 * time.Minute,
 		Workers:           w.riverWorkers,
 		FetchCooldown:     50 * time.Millisecond,
 		FetchPollInterval: 100 * time.Millisecond,
+		Middleware:        middlewares,
 	})
 	if err != nil {
 		return err
@@ -105,6 +111,7 @@ func (w *RiverJobs) initClient() error {
 		return err
 	}
 	return nil
+
 }
 
 func (w *RiverJobs) Use(mwf jobs.JobMiddleware) {
@@ -166,7 +173,7 @@ func (w *RiverJobs) makeRiverJobArgs(job jobs.Job) river.InsertManyParams {
 		}
 	}
 	return river.InsertManyParams{
-		Args:       newRiverJobArgsFrmoJob(job),
+		Args:       newRiverJobArgsFromJob(job),
 		InsertOpts: &insertOpts,
 	}
 }