Merge pull request #95 from btnguyen2k/feat_multitx

Merge feature to dev

README.md

@@ -106,6 +106,7 @@ err = tx.Commit()
 if err != nil {
 	panic(err)
 }
+
 rowsAffected1, err1 := fmt.Println(result1.RowsAffected())
 if err1 != nil {
 	panic(err1)
@@ -123,6 +124,44 @@ fmt.Println("RowsAffected:", rowsAffected2) // output "RowsAffected: 1"
 >
 > You can use [EXISTS function](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ql-functions.exists.html) for condition checking.
 
+Notes on transactions:
+
+- Results of `INSERT`/`UPDATE`/`DELETE` statements are not available until the transaction is committed. Which means, calling
+`RowsAffected()` before `Commit()` will return `0, ErrInTx`.
+- If the connection which has a non-commit/non-rollback transaction is used to execute another statement, the statement is 
+added to the transaction. If the transaction is being committed or rolled back, the execution of the statement will fail
+with error `ErrInTx`. For example:
+
+```go
+conn, _ := db.Conn(context.Background())
+tx, err := conn.BeginTx(context.Background(), nil)
+if err != nil {
+	panic(err)
+}
+result1, _ := tx.Exec(`INSERT INTO "tbltest" VALUE {'app': ?, 'user': ?, 'active': ?}`, "app0", "user1", true)
+
+// the statement is added to the existing transaction
+// also, result2.RowsAffected() is not available until the transaction is committed
+result2, _ := conn.ExecContext(context.Background(), `INSERT INTO "tbltest" VALUE {'app': ?, 'user': ?, 'duration': ?}`, "app0", "user2", 1.23)
+
+err = tx.Commit()
+if err != nil {
+	panic(err)
+}
+
+rowsAffected1, err1 := fmt.Println(result1.RowsAffected())
+if err1 != nil {
+	panic(err1)
+}
+fmt.Println("RowsAffected:", rowsAffected1) // output "RowsAffected: 1"
+
+rowsAffected2, err2 := fmt.Println(result2.RowsAffected())
+if err2 != nil {
+	panic(err2)
+}
+fmt.Println("RowsAffected:", rowsAffected2) // output "RowsAffected: 1"
+```
+
 ## Caveats
 
 **Numerical values** are stored in DynamoDB as floating point numbers. Hence, numbers are always read back as `float64`. 

conn.go

@@ -5,6 +5,7 @@ import (
 	"database/sql/driver"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -13,8 +14,11 @@ import (
 )
 
 var (
-	ErrInTx           = errors.New("statement is in the middle of a transaction, result is not available until committed")
-	ErrInvalidTxStage = errors.New("invalid transaction stage, cannot execute statement ")
+	ErrInTx           = errors.New("there is an ongoing transaction, new transaction/statement or fetching result is not allowed")
+	ErrInvalidTxStage = errors.New("invalid transaction stage")
+	ErrNoTx           = errors.New("no transaction is in progress")
+	ErrTxCommitting   = errors.New("transaction is being committed")
+	ErrTxRollingBack  = errors.New("transaction is being rolled back")
 )
 
 type txMode int
@@ -23,8 +27,10 @@ const (
 	txNone txMode = iota
 	txStarted
 	txCommitting
+	txRollingBack
 )
 
+// txStmt holds a statement to be executed in a transaction.
 type txStmt struct {
 	stmt   *Stmt
 	values []driver.NamedValue
@@ -37,6 +43,8 @@ type executeStatementOutputWrapper func() *dynamodb.ExecuteStatementOutput
 type Conn struct {
 	client     *dynamodb.Client // AWS DynamoDB client
 	timeout    time.Duration
+	lock       sync.Mutex
+	tx         *Tx
 	txMode     txMode
 	txStmtList []*txStmt
 }
@@ -58,15 +66,29 @@ func (c *Conn) ensureContext(ctx context.Context) context.Context {
 }
 
 func (c *Conn) commit() error {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.tx == nil {
+		return ErrNoTx
+	}
+	if c.txMode == txRollingBack {
+		return ErrTxRollingBack
+	}
+	if c.txMode != txStarted && c.txMode != txCommitting {
+		return ErrInvalidTxStage
+	}
+	c.txMode = txCommitting
 	defer func() {
+		c.tx = nil
 		c.txMode = txNone
 		c.txStmtList = nil
 	}()
+
 	if len(c.txStmtList) == 0 {
 		//empty transaction should be successful
 		return nil
 	}
-	c.txMode = txCommitting
+
 	txStmts := make([]types.ParameterizedStatement, len(c.txStmtList))
 	for i, txStmt := range c.txStmtList {
 		params := make([]types.AttributeValue, len(txStmt.values))
@@ -99,7 +121,20 @@ func (c *Conn) commit() error {
 }
 
 func (c *Conn) rollback() error {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.tx == nil {
+		return ErrNoTx
+	}
+	if c.txMode == txCommitting {
+		return ErrTxCommitting
+	}
+	if c.txMode != txStarted && c.txMode != txRollingBack {
+		return ErrInvalidTxStage
+	}
+	c.txMode = txRollingBack
 	defer func() {
+		c.tx = nil
 		c.txMode = txNone
 		c.txStmtList = nil
 	}()
@@ -108,16 +143,24 @@ func (c *Conn) rollback() error {
 
 // execute executes a PartiQL query and returns the result output.
 func (c *Conn) executeContext(ctx context.Context, stmt *Stmt, values []driver.NamedValue) (executeStatementOutputWrapper, error) {
+	//fmt.Printf("[DEBUG] executeContext: in-tx %5v - %s\n", c.tx != nil, stmt.query)
 	if c.txMode == txStarted {
+		// transaction has started and not yet committed or rolled back
+		// --> can add more statements to the transaction
 		txStmt := txStmt{stmt: stmt, values: values}
 		c.txStmtList = append(c.txStmtList, &txStmt)
 		return func() *dynamodb.ExecuteStatementOutput {
 			return txStmt.output
 		}, ErrInTx
 	}
 	if c.txMode != txNone {
+		// transaction is in the middle of committing or rolling back
+		// --> can neither add more statements to the transaction nor execute any statement
 		return nil, ErrInvalidTxStage
 	}
+
+	/* not in transaction mode, execute the statement normally */
+
 	params := make([]types.AttributeValue, len(values))
 	var err error
 	for i, v := range values {
@@ -215,13 +258,19 @@ func (c *Conn) Prepare(query string) (driver.Stmt, error) {
 
 // PrepareContext implements driver.ConnPrepareContext/PrepareContext.
 //
+// Note: since <<VERSION>>, this function returns ErrInTx if there is an outgoing transaction.
+//
 // @Available since v0.2.0
 func (c *Conn) PrepareContext(_ context.Context, query string) (driver.Stmt, error) {
 	return parseQuery(c, query)
 }
 
 // Close implements driver.Conn/Close.
 func (c *Conn) Close() error {
+	if c.tx != nil {
+		//rolling back any outgoing transaction
+		return c.tx.Rollback()
+	}
 	return nil
 }
 
@@ -230,13 +279,19 @@ func (c *Conn) Begin() (driver.Tx, error) {
 	return c.BeginTx(context.Background(), driver.TxOptions{})
 }
 
-// BeginTx implements driver.ConnBeginTx/BeginTx.
+// BeginTx implements driver.Conn/BeginTx.
 //
 // @Available since v0.2.0
 func (c *Conn) BeginTx(_ context.Context, _ driver.TxOptions) (driver.Tx, error) {
-	c.txMode = txStarted
-	c.txStmtList = make([]*txStmt, 0)
-	return &Tx{conn: c}, nil
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.tx == nil {
+		c.tx = &Tx{conn: c}
+		c.txMode = txStarted
+		c.txStmtList = make([]*txStmt, 0)
+		return c.tx, nil
+	}
+	return c.tx, ErrInTx
 }
 
 // CheckNamedValue implements driver.NamedValueChecker/CheckNamedValue.

module_test/stmt_test.go

@@ -245,6 +245,7 @@ func _verifyGSIInfo(t *testing.T, testName string, row map[string]interface{}, g
 }
 
 func _fetchAllRows(dbRows *sql.Rows) ([]map[string]interface{}, error) {
+	defer func() { _ = dbRows.Close() }()
 	colTypes, err := dbRows.ColumnTypes()
 	if err != nil {
 		return nil, err