Skip to content

Commit

Permalink
Add NewFromDB (#34)
Browse files Browse the repository at this point in the history
  • Loading branch information
craigpastro authored Nov 10, 2023
1 parent 6e372dc commit 0f887eb
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 14 deletions.
14 changes: 14 additions & 0 deletions mocks/pgmq.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

25 changes: 11 additions & 14 deletions pgmq.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ type Message struct {
}

type DB interface {
Ping(ctx context.Context) error
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
Expand All @@ -51,31 +52,27 @@ func New(ctx context.Context, connString string) (*PGMQ, error) {
return nil, fmt.Errorf("error creating pool: %w", err)
}

err = pool.Ping(ctx)
if err != nil {
return NewFromDB(ctx, pool)
}

// NewFromDB is a bring your own DB version of New. Given an implementation
// of DB, it will call Ping to ensure the connection has been established,
// then create the PGMQ extension if it does not already exist.
func NewFromDB(ctx context.Context, db DB) (*PGMQ, error) {
if err := db.Ping(ctx); err != nil {
return nil, err
}

_, err = pool.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS pgmq CASCADE")
_, err := db.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS pgmq CASCADE")
if err != nil {
return nil, fmt.Errorf("error creating pgmq extension: %w", err)
}

return &PGMQ{
db: pool,
db: db,
}, nil
}

// MustNew is similar to New, but panics if it encounters an error.
func MustNew(ctx context.Context, connString string) *PGMQ {
q, err := New(ctx, connString)
if err != nil {
panic(err)
}

return q
}

// Close closes the underlying connection pool.
func (p *PGMQ) Close() {
p.db.Close()
Expand Down

0 comments on commit 0f887eb

Please sign in to comment.