Repository: upper/db
Pull request #587 (2020-08-17 21:30:15)

v4: transaction retry


Patch:
diff --git a/adapter/cockroachdb/database.go b/adapter/cockroachdb/database.go
index a9ce271..710fde2 100644
--- a/adapter/cockroachdb/database.go
+++ b/adapter/cockroachdb/database.go
@@ -33,7 +33,7 @@ import (
 	"strings"
 	"time"
 
-	_ "github.com/lib/pq"
+	pq "github.com/lib/pq"
 	db "github.com/upper/db/v4"
 	"github.com/upper/db/v4/internal/sqladapter"
 	"github.com/upper/db/v4/internal/sqladapter/exql"
@@ -48,6 +48,7 @@ func (*database) Template() *exql.Template {
 }
 
 func (*database) OpenDSN(sess sqladapter.Session, dsn string) (*sql.DB, error) {
+	sess.SetMaxTransactionRetries(5)
 	return sql.Open("postgres", dsn)
 }
 
@@ -125,13 +126,19 @@ func (*database) CompileStatement(sess sqladapter.Session, stmt *exql.Statement,
 	return query, args, nil
 }
 
-func (*database) Err(sess sqladapter.Session, err error) error {
+func (*database) Err(err error) error {
 	if err != nil {
 		s := err.Error()
 		// These errors are not exported so we have to check them by they string value.
 		if strings.Contains(s, `too many clients`) || strings.Contains(s, `remaining connection slots are reserved`) || strings.Contains(s, `too many open`) {
 			return db.ErrTooManyClients
 		}
+		if pqErr, ok := err.(*pq.Error); ok {
+			switch pqErr.Code {
+			case "25P02":
+				return db.ErrTransactionAborted
+			}
+		}
 	}
 	return err
 }
diff --git a/errors.go b/errors.go
index ea34e52..8864436 100644
--- a/errors.go
+++ b/errors.go
@@ -57,4 +57,5 @@ var (
 	ErrZeroItemID               = errors.New(`upper: item ID is not defined`)
 	ErrMissingPrimaryKeys       = errors.New(`upper: collection has no primary keys`)
 	ErrWarnSlowQuery            = errors.New(`upper: slow query`)
+	ErrTransactionAborted       = errors.New(`upper: transaction was aborted`)
 )
diff --git a/internal/sqladapter/database.go b/internal/sqladapter/database.go
index 6c5618f..d6ba201 100644
--- a/internal/sqladapter/database.go
+++ b/internal/sqladapter/database.go
@@ -159,9 +159,9 @@ type Session interface {
 
 	NewTx(ctx context.Context) (sqlbuilder.Tx, error)
 
-	Tx(fn func(sess sqlbuilder.Tx) error) error
+	Tx(fn func(sess db.Session) error) error
 
-	TxContext(ctx context.Context, fn func(sess sqlbuilder.Tx) error) error
+	TxContext(ctx context.Context, fn func(sess db.Session) error) error
 
 	WithContext(context.Context) sqlbuilder.Session
 
@@ -238,16 +238,16 @@ func (sess *session) WithContext(ctx context.Context) sqlbuilder.Session {
 	return newDB
 }
 
-func (sess *session) Tx(fn func(tx sqlbuilder.Tx) error) error {
+func (sess *session) Tx(fn func(sess db.Session) error) error {
 	return TxContext(sess.Context(), sess, fn)
 }
 
-func (sess *session) TxContext(ctx context.Context, fn func(sess sqlbuilder.Tx) error) error {
+func (sess *session) TxContext(ctx context.Context, fn func(sess db.Session) error) error {
 	return TxContext(ctx, sess, fn)
 }
 
 func (sess *session) Err(errIn error) (errOur error) {
-	if convertError, ok := errIn.(errorConverter); ok {
+	if convertError, ok := sess.adapter.(errorConverter); ok {
 		return convertError.Err(errIn)
 	}
 	return errIn
@@ -630,7 +630,6 @@ func (sess *session) NewClone(adapter AdapterSession, checkConn bool) (Session,
 }
 
 func (sess *session) Close() error {
-
 	defer func() {
 		sess.sqlDBMu.Lock()
 		sess.sqlDB = nil
diff --git a/internal/sqladapter/tx.go b/internal/sqladapter/tx.go
index 817a671..1cc4f94 100644
--- a/internal/sqladapter/tx.go
+++ b/internal/sqladapter/tx.go
@@ -24,7 +24,10 @@ package sqladapter
 import (
 	"context"
 	"database/sql"
+	"errors"
+	"fmt"
 	"sync/atomic"
+	"time"
 
 	db "github.com/upper/db/v4"
 	"github.com/upper/db/v4/sqlbuilder"
@@ -94,18 +97,39 @@ func (w *sessionTx) Rollback() error {
 }
 
 // TxContext creates a transaction context and runs fn within it.
-func TxContext(ctx context.Context, sess sqlbuilder.Session, fn func(tx sqlbuilder.Tx) error) error {
-	tx, err := sess.NewTx(ctx)
-	if err != nil {
-		return err
+func TxContext(ctx context.Context, sess sqlbuilder.Session, fn func(tx db.Session) error) error {
+
+	txFn := func(sess sqlbuilder.Session) error {
+		tx, err := sess.NewTx(ctx)
+		if err != nil {
+			return err
+		}
+		defer tx.Close()
+
+		if err := fn(tx); err != nil {
+			if rollbackErr := tx.Rollback(); rollbackErr != nil {
+				return fmt.Errorf("%s: %w", rollbackErr, err)
+			}
+			return err
+		}
+		return tx.Commit()
 	}
 
-	defer tx.Close()
-	if err := fn(tx); err != nil {
-		_ = tx.Rollback()
-		return err
+	var txErr error
+	for i := 0; i < sess.MaxTransactionRetries(); i++ {
+		txErr = sess.(*session).Err(txFn(sess))
+		if txErr == nil {
+			return nil
+		}
+		if errors.Is(txErr, db.ErrTransactionAborted) {
+			time.Sleep(time.Millisecond * 100 * time.Duration(i))
+			continue
+		}
+		return txErr
 	}
-	return tx.Commit()
+
+	return fmt.Errorf("db: giving up trying to commit transaction: %w", txErr)
+
 }
 
 var (
@@ -119,9 +143,10 @@ type txWrapper struct {
 
 var (
 	_ = sqlbuilder.Tx(&txWrapper{})
+	_ = sqlbuilder.Session(&txWrapper{})
 )
 
-func (t *txWrapper) WithContext(ctx context.Context) sqlbuilder.Tx {
+func (t *txWrapper) WithContext(ctx context.Context) sqlbuilder.Session {
 	newTx := *t
 	newTx.SessionTx.SetContext(ctx)
 	return &newTx
diff --git a/settings.go b/settings.go
index b54714a..98811f0 100644
--- a/settings.go
+++ b/settings.go
@@ -60,6 +60,14 @@ type Settings interface {
 	// MaxOpenConns returns the default maximum number of open connections to the
 	// database.
 	MaxOpenConns() int
+
+	// SetMaxTransactionRetries sets the number of times a transaction can
+	// be retried.
+	SetMaxTransactionRetries(int)
+
+	// MaxTransactionRetries returns the maximum number of times a
+	// transaction can be retried.
+	MaxTransactionRetries() int
 }
 
 type settings struct {
@@ -70,6 +78,8 @@ type settings struct {
 	connMaxLifetime time.Duration
 	maxOpenConns    int
 	maxIdleConns    int
+
+	maxTransactionRetries int
 }
 
 func (c *settings) binaryOption(opt *uint32) bool {
@@ -116,6 +126,21 @@ func (c *settings) MaxIdleConns() int {
 	return c.maxIdleConns
 }
 
+func (c *settings) SetMaxTransactionRetries(n int) {
+	c.Lock()
+	c.maxTransactionRetries = n
+	c.Unlock()
+}
+
+func (c *settings) MaxTransactionRetries() int {
+	c.RLock()
+	defer c.RUnlock()
+	if c.maxTransactionRetries < 1 {
+		return 1
+	}
+	return c.maxTransactionRetries
+}
+
 func (c *settings) SetMaxOpenConns(n int) {
 	c.Lock()
 	c.maxOpenConns = n
@@ -137,6 +162,7 @@ func NewSettings() Settings {
 		connMaxLifetime:               def.connMaxLifetime,
 		maxIdleConns:                  def.maxIdleConns,
 		maxOpenConns:                  def.maxOpenConns,
+		maxTransactionRetries:         def.maxTransactionRetries,
 	}
 }
 
@@ -147,4 +173,5 @@ var DefaultSettings Settings = &settings{
 	connMaxLifetime:               time.Duration(0),
 	maxIdleConns:                  10,
 	maxOpenConns:                  0,
+	maxTransactionRetries:         1,
 }
diff --git a/sqlbuilder/wrapper.go b/sqlbuilder/wrapper.go
index efef4a1..3c5ae59 100644
--- a/sqlbuilder/wrapper.go
+++ b/sqlbuilder/wrapper.go
@@ -54,7 +54,7 @@ type Tx interface {
 	// WithContext returns a copy of the transaction that uses the given context
 	// as default. Copies are safe to use concurrently but they're backed by the
 	// same *sql.Tx, so any copy may commit or rollback the parent transaction.
-	WithContext(context.Context) Tx
+	// WithContext(context.Context) Tx
 
 	// SetTxOptions sets the default TxOptions that is going to be used for new
 	// transactions created in the session.
@@ -82,13 +82,13 @@ type Session interface {
 	// it to the function fn. If fn returns no error the transaction is commited,
 	// else the transaction is rolled back. After being commited or rolled back
 	// the transaction is closed automatically.
-	Tx(fn func(sess Tx) error) error
+	Tx(fn func(sess db.Session) error) error
 
 	// TxContext creates a transaction block on the given context and passes it to
 	// the function fn. If fn returns no error the transaction is commited, else
 	// the transaction is rolled back. After being commited or rolled back the
 	// transaction is closed automatically.
-	TxContext(ctx context.Context, fn func(sess Tx) error) error
+	TxContext(ctx context.Context, fn func(sess db.Session) error) error
 
 	// Context returns the context used as default for queries on this session
 	// and for new transactions.  If no context has been set, a default


Test patch:
diff --git a/adapter/cockroachdb/cockroachdb_test.go b/adapter/cockroachdb/cockroachdb_test.go
index 59390fa..025c006 100644
--- a/adapter/cockroachdb/cockroachdb_test.go
+++ b/adapter/cockroachdb/cockroachdb_test.go
@@ -159,7 +159,7 @@ func (s *AdapterTests) SkipTest_Issue469_BadConnection() {
 
 	// At this point the server should have disconnected us. Let's try to create
 	// a transaction anyway.
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		_, err = sess.Collection("artist").Find().Count()
@@ -175,7 +175,7 @@ func (s *AdapterTests) SkipTest_Issue469_BadConnection() {
 	_, err = sess.Exec(`SET statement_timeout=1000`)
 	s.NoError(err)
 
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		// This query should succeed.
diff --git a/adapter/mysql/mysql_test.go b/adapter/mysql/mysql_test.go
index 58a2164..4204c25 100644
--- a/adapter/mysql/mysql_test.go
+++ b/adapter/mysql/mysql_test.go
@@ -31,8 +31,8 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/suite"
+	"github.com/upper/db/v4"
 	"github.com/upper/db/v4/internal/testsuite"
-	"github.com/upper/db/v4/sqlbuilder"
 )
 
 type int64Compat int64
@@ -162,7 +162,7 @@ func (s *AdapterTests) TestIssue469_BadConnection() {
 
 	// At this point the server should have disconnected us. Let's try to create
 	// a transaction anyway.
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		_, err = sess.Collection("artist").Find().Count()
@@ -178,7 +178,7 @@ func (s *AdapterTests) TestIssue469_BadConnection() {
 	_, err = sess.Exec(`SET SESSION wait_timeout=1`)
 	s.NoError(err)
 
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		// This query should succeed.
diff --git a/adapter/postgresql/postgresql_test.go b/adapter/postgresql/postgresql_test.go
index 7f40c38..56f33ce 100644
--- a/adapter/postgresql/postgresql_test.go
+++ b/adapter/postgresql/postgresql_test.go
@@ -157,7 +157,7 @@ func (s *AdapterTests) Test_Issue469_BadConnection() {
 
 	// At this point the server should have disconnected us. Let's try to create
 	// a transaction anyway.
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		_, err = sess.Collection("artist").Find().Count()
@@ -173,7 +173,7 @@ func (s *AdapterTests) Test_Issue469_BadConnection() {
 	_, err = sess.Exec(`SET SESSION idle_in_transaction_session_timeout=1000`)
 	s.NoError(err)
 
-	err = sess.Tx(func(sess sqlbuilder.Tx) error {
+	err = sess.Tx(func(sess db.Session) error {
 		var err error
 
 		// This query should succeed.
diff --git a/internal/testsuite/record_suite.go b/internal/testsuite/record_suite.go
index 19fd46f..5fb3a60 100644
--- a/internal/testsuite/record_suite.go
+++ b/internal/testsuite/record_suite.go
@@ -76,7 +76,7 @@ func (s *RecordTestSuite) BeforeTest(suiteName, testName string) {
 	err := s.TearUp()
 	s.NoError(err)
 
-	sess := s.Helper.Session().(sqlbuilder.Session)
+	sess := s.Helper.Session()
 
 	cols, err := sess.Collections()
 	s.NoError(err)
@@ -235,7 +235,7 @@ func (s *RecordTestSuite) TestSelectOnlyIDs() {
 }
 
 func (s *RecordTestSuite) TestTx() {
-	sess := s.Session().(sqlbuilder.Session)
+	sess := s.SQLBuilder()
 
 	user := User{Username: "peter"}
 	err := sess.Save(&user)
@@ -243,21 +243,21 @@ func (s *RecordTestSuite) TestTx() {
 
 	// This transaction should fail because user is a UNIQUE value and we already
 	// have a "peter".
-	err = sess.Tx(func(tx sqlbuilder.Tx) error {
+	err = sess.Tx(func(tx db.Session) error {
 		return tx.Save(&User{Username: "peter"})
 	})
 	s.Error(err)
 
 	// This transaction should fail because user is a UNIQUE value and we already
 	// have a "peter".
-	err = sess.Tx(func(tx sqlbuilder.Tx) error {
+	err = sess.Tx(func(tx db.Session) error {
 		return tx.Save(&User{Username: "peter"})
 	})
 	s.Error(err)
 
 	// This transaction will have no errors, but we'll produce one in order for
 	// it to rollback at the last moment.
-	err = sess.Tx(func(tx sqlbuilder.Tx) error {
+	err = sess.Tx(func(tx db.Session) error {
 		if err := tx.Save(&User{Username: "Joe"}); err != nil {
 			return err
 		}
@@ -272,7 +272,7 @@ func (s *RecordTestSuite) TestTx() {
 
 	// Attempt to add two new unique values, if the transaction above had not
 	// been rolled back this transaction will fail.
-	err = sess.Tx(func(tx sqlbuilder.Tx) error {
+	err = sess.Tx(func(tx db.Session) error {
 		if err := tx.Save(&User{Username: "Joe"}); err != nil {
 			return err
 		}
@@ -286,7 +286,7 @@ func (s *RecordTestSuite) TestTx() {
 	s.NoError(err)
 
 	// If the transaction above was successful, this one will fail.
-	err = sess.Tx(func(tx sqlbuilder.Tx) error {
+	err = sess.Tx(func(tx db.Session) error {
 		if err := tx.Save(&User{Username: "Joe"}); err != nil {
 			return err
 		}