streamline the database interfaces and add pgx adapter (#29)

Author: julesjcraske

.github/workflows/pr.yml

@@ -1,5 +1,6 @@
 name: pr
 on:
+  pull_request:
   push:
     branches-ignore:
       - master

docker-compose.yml

@@ -1,7 +1,6 @@
 # for Watermill development purposes.
 # For Watermill based application docker please check https://watermill.io/docs/getting-started/
 
-version: '3'
 services:
   mysql:
     image: mysql:8.0

go.mod

@@ -1,15 +1,16 @@
 module github.com/ThreeDotsLabs/watermill-sql/v4
 
-go 1.21
+go 1.22
 
-toolchain go1.23.0
+toolchain go1.23.3
 
 require (
 	github.com/ThreeDotsLabs/watermill v1.4.0-rc.2
 	github.com/go-sql-driver/mysql v1.4.1
-	github.com/jackc/pgx/v4 v4.18.2
-	github.com/lib/pq v1.10.2
+	github.com/jackc/pgx/v5 v5.7.1
+	github.com/lib/pq v1.10.9
 	github.com/oklog/ulid v1.3.1
+	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.9.0
 )
 
@@ -19,18 +20,15 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
-	github.com/jackc/chunkreader/v2 v2.0.1 // indirect
-	github.com/jackc/pgconn v1.14.3 // indirect
-	github.com/jackc/pgio v1.0.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
-	github.com/jackc/pgproto3/v2 v2.3.3 // indirect
-	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
-	github.com/jackc/pgtype v1.14.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/lithammer/shortuuid/v3 v3.0.7 // indirect
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/rogpeppe/go-internal v1.13.1 // indirect
 	github.com/sony/gobreaker v1.0.0 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
+	golang.org/x/sync v0.10.0 // indirect
 	golang.org/x/text v0.21.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

go.sum

@@ -0,0 +1,127 @@
+package sql
+
+import (
+	"context"
+	stdSQL "database/sql"
+	"fmt"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+)
+
+type Conn interface {
+	BeginTx(ctx context.Context, options pgx.TxOptions) (pgx.Tx, error)
+	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
+	Query(ctx context.Context, sql string, arguments ...any) (pgx.Rows, error)
+}
+
+type PgxBeginner struct {
+	Conn
+}
+
+func BeginnerFromPgx(conn Conn) Beginner {
+	return PgxBeginner{conn}
+}
+
+type PgxTx struct {
+	pgx.Tx
+	ctx context.Context
+}
+
+type PgxResult struct {
+	pgconn.CommandTag
+}
+
+type PgxRows struct {
+	pgx.Rows
+}
+
+func (c PgxBeginner) BeginTx(ctx context.Context, options *stdSQL.TxOptions) (Tx, error) {
+	opts := pgx.TxOptions{}
+	if options != nil {
+		iso, err := toPgxIsolationLevel(options.Isolation)
+		if err != nil {
+			return nil, err
+		}
+
+		opts = pgx.TxOptions{
+			IsoLevel:       iso,
+			AccessMode:     toPgxAccessMode(options.ReadOnly),
+			DeferrableMode: "",
+		}
+	}
+
+	tx, err := c.Conn.BeginTx(ctx, opts)
+
+	return &PgxTx{
+		Tx:  tx,
+		ctx: ctx,
+	}, err
+}
+
+func (c PgxBeginner) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
+	res, err := c.Conn.Exec(ctx, query, args...)
+
+	return PgxResult{res}, err
+}
+
+func (c PgxBeginner) QueryContext(ctx context.Context, query string, args ...any) (Rows, error) {
+	rows, err := c.Conn.Query(ctx, query, args...)
+
+	return PgxRows{rows}, err
+}
+
+func (t PgxTx) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
+	res, err := t.Tx.Exec(ctx, query, args...)
+
+	return PgxResult{res}, err
+}
+
+func (t PgxTx) QueryContext(ctx context.Context, query string, args ...any) (Rows, error) {
+	rows, err := t.Tx.Query(ctx, query, args...)
+
+	return PgxRows{rows}, err
+}
+
+func (t PgxTx) Rollback() error {
+	return t.Tx.Rollback(context.WithoutCancel(t.ctx))
+}
+
+func (t PgxTx) Commit() error {
+	return t.Tx.Commit(t.ctx)
+}
+
+func (p PgxResult) RowsAffected() (int64, error) {
+	return p.CommandTag.RowsAffected(), nil
+}
+
+func (p PgxRows) Close() error {
+	p.Rows.Close()
+
+	return nil
+}
+
+func toPgxIsolationLevel(level stdSQL.IsolationLevel) (pgx.TxIsoLevel, error) {
+	switch level {
+	case stdSQL.LevelReadUncommitted:
+		return pgx.ReadUncommitted, nil
+	case stdSQL.LevelReadCommitted:
+		return pgx.ReadCommitted, nil
+	case stdSQL.LevelRepeatableRead:
+		return pgx.RepeatableRead, nil
+	case stdSQL.LevelSerializable:
+		return pgx.Serializable, nil
+	case stdSQL.LevelSnapshot:
+		return pgx.Serializable, fmt.Errorf("pgx does not support snapshot isolation")
+	default:
+		return pgx.Serializable, nil
+	}
+}
+
+func toPgxAccessMode(readOnly bool) pgx.TxAccessMode {
+	if readOnly {
+		return pgx.ReadOnly
+	}
+
+	return pgx.ReadWrite
+}

pkg/sql/adapters_pgx.go

@@ -0,0 +1,49 @@
+package sql
+
+import (
+	"context"
+	"database/sql"
+)
+
+type StdSQLBeginner struct {
+	SQLBeginner
+}
+
+func BeginnerFromStdSQL(sqlBeginner SQLBeginner) Beginner {
+	return StdSQLBeginner{sqlBeginner}
+}
+
+type StdSQLTx struct {
+	*sql.Tx
+}
+
+func TxFromStdSQL(tx *sql.Tx) Tx {
+	return &StdSQLTx{tx}
+}
+
+// BeginTx converts the stdSQL.Tx struct to our Tx interface
+func (c StdSQLBeginner) BeginTx(ctx context.Context, options *sql.TxOptions) (Tx, error) {
+	tx, err := c.SQLBeginner.BeginTx(ctx, options)
+
+	return &StdSQLTx{tx}, err
+}
+
+// ExecContext converts the stdSQL.Result struct to our Result interface
+func (c StdSQLBeginner) ExecContext(ctx context.Context, query string, args ...interface{}) (Result, error) {
+	return c.SQLBeginner.ExecContext(ctx, query, args...)
+}
+
+// QueryContext converts the stdSQL.Rows struct to our Rows interface
+func (c StdSQLBeginner) QueryContext(ctx context.Context, query string, args ...interface{}) (Rows, error) {
+	return c.SQLBeginner.QueryContext(ctx, query, args...)
+}
+
+// ExecContext converts the stdSQL.Result struct to our Result interface
+func (t StdSQLTx) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
+	return t.Tx.ExecContext(ctx, query, args...)
+}
+
+// QueryContext converts the stdSQL.Rows struct to our Rows interface
+func (t StdSQLTx) QueryContext(ctx context.Context, query string, args ...any) (Rows, error) {
+	return t.Tx.QueryContext(ctx, query, args...)
+}

pkg/sql/adapters_std_sql.go

@@ -2,7 +2,6 @@ package sql
 
 import (
 	"context"
-	"database/sql"
 )
 
 type contextKey string
@@ -11,7 +10,7 @@ const (
 	txContextKey contextKey = "tx"
 )
 
-func setTxToContext(ctx context.Context, tx *sql.Tx) context.Context {
+func setTxToContext(ctx context.Context, tx Tx) context.Context {
 	return context.WithValue(ctx, txContextKey, tx)
 }
 
@@ -21,7 +20,7 @@ func setTxToContext(ctx context.Context, tx *sql.Tx) context.Context {
 //
 // It is useful when you want to ensure that data is updated only when the message is processed.
 // Example usage: https://github.com/ThreeDotsLabs/watermill/tree/master/_examples/real-world-examples/exactly-once-delivery-counter
-func TxFromContext(ctx context.Context) (*sql.Tx, bool) {
-	tx, ok := ctx.Value(txContextKey).(*sql.Tx)
+func TxFromContext(ctx context.Context) (Tx, bool) {
+	tx, ok := ctx.Value(txContextKey).(Tx)
 	return tx, ok
 }

pkg/sql/context.go

@@ -1,7 +1,6 @@
 package sql
 
 import (
-	"database/sql"
 	"fmt"
 	"strings"
 
@@ -28,7 +27,7 @@ func (c *DelayedPostgreSQLPublisherConfig) setDefaults() {
 
 // NewDelayedPostgreSQLPublisher creates a new Publisher that stores messages in PostgreSQL with a delay.
 // The delay can be set per message with the Watermill's components/delay metadata.
-func NewDelayedPostgreSQLPublisher(db *sql.DB, config DelayedPostgreSQLPublisherConfig) (message.Publisher, error) {
+func NewDelayedPostgreSQLPublisher(db ContextExecutor, config DelayedPostgreSQLPublisherConfig) (message.Publisher, error) {
 	config.setDefaults()
 
 	publisherConfig := PublisherConfig{
@@ -82,7 +81,7 @@ func (c *DelayedPostgreSQLSubscriberConfig) setDefaults() {
 
 // NewDelayedPostgreSQLSubscriber creates a new Subscriber that reads messages from PostgreSQL with a delay.
 // The delay can be set per message with the Watermill's components/delay metadata.
-func NewDelayedPostgreSQLSubscriber(db *sql.DB, config DelayedPostgreSQLSubscriberConfig) (message.Subscriber, error) {
+func NewDelayedPostgreSQLSubscriber(db Beginner, config DelayedPostgreSQLSubscriberConfig) (message.Subscriber, error) {
 	config.setDefaults()
 
 	where := fmt.Sprintf("(metadata->>'%v')::timestamptz < NOW() AT TIME ZONE 'UTC'", delay.DelayedUntilKey)

pkg/sql/delayed_postgresql.go

@@ -2,7 +2,6 @@ package sql
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"time"
 
@@ -35,7 +34,7 @@ func (q DelayedRequeuer) Run(ctx context.Context) error {
 // DelayedRequeuerConfig is a configuration for DelayedRequeuer.
 type DelayedRequeuerConfig struct {
 	// DB is a database connection. Required.
-	DB *sql.DB
+	DB Beginner
 
 	// Publisher is a publisher that will be used to publish requeued messages. Required.
 	Publisher message.Publisher

pkg/sql/delayed_requeuer.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/ThreeDotsLabs/watermill"
 	"github.com/ThreeDotsLabs/watermill/message"
@@ -133,8 +134,11 @@ func (p *Publisher) initializeSchema(topic string) error {
 		return nil
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*60)
+	defer cancel()
+
 	if err := initializeSchema(
-		context.Background(),
+		ctx,
 		topic,
 		p.logger,
 		p.db,