Pass ctx to constructor (#14)

README.md

@@ -30,7 +30,7 @@ import (
 )
 
 func main() {
-    q, err := pgmq.New("postgres://postgres:password@localhost:5432/postgres")
+    q, err := pgmq.New(context.Background(), "postgres://postgres:password@localhost:5432/postgres")
     if err != nil {
         panic(err)
     }

pgmq.go

@@ -4,18 +4,19 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"log"
 	"time"
 
-	"github.com/craigpastro/retrier"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 const vtDefault = 30
 
-var ErrNoRows = errors.New("pgmq: no rows in result set")
+var (
+	ErrNoRows = errors.New("pgmq: no rows in result set")
+	ErrPing   = errors.New("pgmq: failed to ping db")
+)
 
 type Message struct {
 	MsgID      int64
@@ -38,31 +39,25 @@ type PGMQ struct {
 	db DB
 }
 
-// New establishes a connection to Postgres given by the connString, then
-// creates the pgmq extension if it does not already exist.
-func New(connString string) (*PGMQ, error) {
+// New establishes a connection to Postgres given by the connString, checks connection, if check is failed,
+// returns ErrPing, that can be retried, then creates the pgmq extension if it does not already exist.
+func New(ctx context.Context, connString string) (*PGMQ, error) {
 	cfg, err := pgxpool.ParseConfig(connString)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing connection string: %w", err)
 	}
 
-	pool, err := pgxpool.NewWithConfig(context.Background(), cfg)
+	pool, err := pgxpool.NewWithConfig(ctx, cfg)
 	if err != nil {
 		return nil, fmt.Errorf("error creating pool: %w", err)
 	}
 
-	err = retrier.Do(func() error {
-		if err = pool.Ping(context.Background()); err != nil {
-			log.Println("waiting for Postgres")
-			return err
-		}
-		return nil
-	}, retrier.NewExponentialBackoff())
+	err = pool.Ping(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("error connecting to Postgres: %w", err)
+		return nil, errors.Join(err, ErrPing)
 	}
 
-	_, err = pool.Exec(context.Background(), "create extension if not exists pgmq cascade")
+	_, err = pool.Exec(ctx, "create extension if not exists pgmq cascade")
 	if err != nil {
 		return nil, fmt.Errorf("error creating pgmq extension: %w", err)
 	}
@@ -73,8 +68,8 @@ func New(connString string) (*PGMQ, error) {
 }
 
 // MustNew is similar to New, but panics if it encounters an error.
-func MustNew(connString string) *PGMQ {
-	q, err := New(connString)
+func MustNew(ctx context.Context, connString string) *PGMQ {
+	q, err := New(ctx, connString)
 	if err != nil {
 		panic(err)
 	}

pgmq_test.go

@@ -8,12 +8,13 @@ import (
 	"testing"
 	"time"
 
-	"github.com/craigpastro/pgmq-go/mocks"
+	"github.com/craigpastro/retrier"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/stretchr/testify/require"
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
 	"go.uber.org/mock/gomock"
+	"github.com/craigpastro/pgmq-go/mocks"
 )
 
 var q *PGMQ
@@ -53,7 +54,12 @@ func TestMain(m *testing.M) {
 
 	connString := fmt.Sprintf("postgres://postgres:password@%s:%s/postgres", host, port.Port())
 
-	q = MustNew(connString)
+	q, err = retrier.DoWithData(func() (*PGMQ, error) {
+		return New(ctx, connString)
+	}, retrier.NewExponentialBackoff())
+	if err != nil {
+		panic(err)
+	}
 
 	code := m.Run()