DB connection with transaction

pg.go

package connection

import (
	"context"
	"database/sql"

	"github.com/pkg/errors"
)

// PG is a connection wrapper with transaction decorators
type PG struct {
	sql.DB
}

type ctxTxKey struct{}

// WithTx execute callback with transaction in context
func (pg *PG) WithTx(ctx context.Context, fn func(ctx context.Context) error) (err error) {
	tx, alreadyHasTx := ctx.Value(ctxTxKey{}).(*sql.Tx)
	if !alreadyHasTx {
		tx, err = pg.BeginTx(ctx, nil)
		if err != nil {

			return errors.WithStack(err)
		}
		ctx = context.WithValue(ctx, ctxTxKey{}, tx)
	}

	err = errors.WithStack(fn(ctx))

	if alreadyHasTx {

		return err
	}
	if err == nil {

		return errors.WithStack(tx.Commit())
	}

	tx.Rollback()

	return err
}

// ExtractTx from context or create new
func (pg *PG) ExtractTx(ctx context.Context, fn func(context.Context, *sql.Tx) error) error {

	return pg.WithTx(ctx, func(ctx context.Context) error {
		tx := ctx.Value(ctxTxKey{}).(*sql.Tx)

		return errors.WithStack(fn(ctx, tx))
	})
}