Add a FlexibleTransaction type and constructor (#61)

* Add a FlexibleTransaction type and constructor

This allows to avoid the boilerplate of handling transactions between functions by wrapping DB into a state tracking shim that simulates a Tx.

Author: perrito666

db/chain/chain.go

@@ -28,8 +28,8 @@ func NewExpressionChain(db connection.DB) *ExpressionChain {
 	return &ExpressionChain{db: db}
 }
 
-// NewNoDB creates an expression chain withouth the db, mostly with the purpose of making a more
-// abbreviated syntax for transient ExpresionChains such as CTE or subquery ones.
+// NewNoDB creates an expression chain without the db, mostly with the purpose of making a more
+// abbreviated syntax for transient ExpressionChains such as CTE or sub-query ones.
 func NewNoDB() *ExpressionChain {
 	return &ExpressionChain{}
 }

db/connection/connection.go

@@ -16,9 +16,11 @@ package connection
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ShiftLeftSecurity/gaum/v2/db/logging"
@@ -117,6 +119,97 @@ type DB interface {
 	BulkInsert(ctx context.Context, tableName string, columns []string, values [][]interface{}) (execError error)
 }
 
+var _ DB = (*FlexibleTransaction)(nil)
+
+// FlexibleTransaction allows for a DB transaction to be passed through functions and avoid multiple commit/rollbacks
+// it also takes care of some of the most repeated checks at the time of commit/rollback and tx checking.
+type FlexibleTransaction struct {
+	DB
+	rolled               bool
+	concurrencySafeguard sync.Mutex
+}
+
+// Cleanup is an implementation of TXFinishFunc for FlexibleTransaction, it handles an attempt to either Commit
+// or rollback a transaction depending on the perceived outcome: If someone invoked rollback on the FlexibleTransaction
+// we assume the process went wrong and will rollback all. This is intended as a way to mitigate the lack of different
+// abstractions for Transaction and Connection in the current version of gaum, retaining the ability to finalize the
+// transaction at the initiator level.
+// This does however allow some bad habits such as functions acting different depending on if they think they receive
+// a transaction or a connection instead of having two functions that force the former or later as arguments.
+func (f *FlexibleTransaction) Cleanup(ctx context.Context) (bool, bool, error) {
+	f.concurrencySafeguard.Lock()
+	defer f.concurrencySafeguard.Unlock()
+	if f.DB == nil {
+		return false, false, nil
+	}
+	if f.rolled {
+		if err := f.DB.RollbackTransaction(ctx); err != nil {
+			return false, false, fmt.Errorf("rolling back transaction: %w", err)
+		}
+		return false, true, nil
+	}
+
+	if err := f.DB.CommitTransaction(ctx); err != nil {
+		return false, false, fmt.Errorf("committing transaction: %w", err)
+	}
+	return true, false, nil
+}
+
+// TXFinishFunc represents an all-encompassing function that either rolls or commits a tx based on the outcome.
+type TXFinishFunc func(ctx context.Context) (committed, rolled bool, err error)
+
+// BeginTransaction will wrap the passed DB into a transaction handler that supports it being used with less care
+// and prevents having to check if we are already in a tx and failures due to eager committers.
+func BeginTransaction(ctx context.Context, conn DB) (DB, TXFinishFunc, error) {
+	// this can happen so let's work around it
+	ft, isFT := conn.(*FlexibleTransaction)
+	if isFT {
+		return ft, func(ctx2 context.Context) (bool, bool, error) {
+			return false, false, nil
+		}, nil
+	}
+
+	// the underlying conn is a tx, let's be careful not to commit/rollback it
+	if conn.IsTransaction() {
+		return &FlexibleTransaction{
+				DB: conn,
+			},
+			func(ctx2 context.Context) (bool, bool, error) {
+				return false, false, nil
+			},
+			nil
+
+	}
+
+	tx, err := conn.BeginTransaction(ctx)
+	if err != nil {
+		return nil, nil, fmt.Errorf("beginning transaction: %w", err)
+	}
+
+	f := &FlexibleTransaction{
+		DB: tx,
+	}
+	return f, f.Cleanup, nil
+}
+
+// BeginTransaction implements DB for FlexibleTransaction
+func (f *FlexibleTransaction) BeginTransaction(ctx context.Context) (DB, error) {
+	return f, nil
+}
+
+// CommitTransaction implements DB for FlexibleTransaction
+func (f *FlexibleTransaction) CommitTransaction(ctx context.Context) error {
+	return nil
+}
+
+// RollbackTransaction implements DB for FlexibleTransaction
+func (f *FlexibleTransaction) RollbackTransaction(ctx context.Context) error {
+	f.concurrencySafeguard.Lock()
+	defer f.concurrencySafeguard.Unlock()
+	f.rolled = true
+	return nil
+}
+
 // EscapeArgs return the query and args with the argument placeholder escaped.
 func EscapeArgs(query string, args []interface{}) (string, []interface{}, error) {
 	// TODO: make this a bit less ugly

db/connection/connection_test.go

@@ -0,0 +1,190 @@
+package connection
+
+import (
+	"context"
+	"testing"
+)
+
+type fakeConn struct {
+	DB
+	begin    int
+	commit   int
+	rollback int
+	isTx     bool
+}
+
+func (f *fakeConn) BeginTransaction(ctx context.Context) (DB, error) {
+	f.begin++
+	f.isTx = true
+	return f, nil
+}
+
+func (f *fakeConn) CommitTransaction(ctx context.Context) error {
+	f.commit++
+	return nil
+}
+
+func (f *fakeConn) RollbackTransaction(ctx context.Context) error {
+	f.rollback++
+	return nil
+}
+
+func (f *fakeConn) IsTransaction() bool {
+	return f.isTx
+}
+
+var _ DB = (*fakeConn)(nil)
+
+func TestFlexibleTransactionSucceeds(t *testing.T) {
+	// Multiple TX begins, multiple commits
+	fc := &fakeConn{}
+	ctx := context.Background()
+	tx, cleanup, err := BeginTransaction(ctx, fc)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 0; i < 10; i++ {
+		if err := tx.CommitTransaction(ctx); err != nil {
+			t.Logf("Repetitive commit N %d", i+1)
+			t.Fatal(err)
+		}
+	}
+	committed, rolledBack, err := cleanup(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !committed {
+		t.Log("tx was not committed but we expected it to")
+		t.FailNow()
+	}
+	if rolledBack {
+		t.Log("tx was rolled back and we did not expect it to")
+		t.FailNow()
+	}
+
+	if fc.begin != 1 {
+		t.Logf("begin was called %d times in the underlying conn but we expected 1", fc.begin)
+		t.FailNow()
+	}
+
+	if fc.commit != 1 {
+		t.Logf("commit was called %d times in the underlying conn but we expected 1", fc.commit)
+		t.FailNow()
+	}
+
+	if fc.rollback != 0 {
+		t.Logf("rollback was called %d times in the underlying conn but we expected 0", fc.rollback)
+		t.FailNow()
+	}
+}
+
+func TestFlexibleTransactionRollbackTransaction(t *testing.T) {
+	// Multiple TX begins, multiple commits
+	fc := &fakeConn{}
+	ctx := context.Background()
+	tx, cleanup, err := BeginTransaction(ctx, fc)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for i := 0; i < 10; i++ {
+		if err := tx.CommitTransaction(ctx); err != nil {
+			t.Logf("Repetitive commit N %d", i+1)
+			t.Fatal(err)
+		}
+	}
+	if err := tx.RollbackTransaction(ctx); err != nil {
+		t.Fatal(err)
+	}
+	committed, rolledBack, err := cleanup(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if committed {
+		t.Log("tx was committed but we expected it not to")
+		t.FailNow()
+	}
+	if !rolledBack {
+		t.Log("tx was not rolled back and we expected it to")
+		t.FailNow()
+	}
+
+	if fc.begin != 1 {
+		t.Logf("begin was called %d times in the underlying conn but we expected 1", fc.begin)
+		t.FailNow()
+	}
+
+	if fc.commit != 0 {
+		t.Logf("commit was called %d times in the underlying conn but we expected 0", fc.commit)
+		t.FailNow()
+	}
+
+	if fc.rollback != 1 {
+		t.Logf("rollback was called %d times in the underlying conn but we expected 1", fc.rollback)
+		t.FailNow()
+	}
+}
+
+func TestFlexibleTransactionRecursive(t *testing.T) {
+	// Multiple TX begins, multiple commits
+	fc := &fakeConn{}
+	ctx := context.Background()
+	tx, cleanup, err := BeginTransaction(ctx, fc)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tx, innerCleanup, err := BeginTransaction(ctx, fc)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// we call it early to see if it really is a noop
+	innerCommit, innerRollback, err := innerCleanup(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if innerCommit {
+		t.Log("commit should not have happened on inner cleanup function")
+		t.FailNow()
+	}
+	if innerRollback {
+		t.Log("rollback should not have happened on inner cleanup function")
+		t.FailNow()
+	}
+	// notice that we are using the inner tx
+	for i := 0; i < 10; i++ {
+		if err := tx.CommitTransaction(ctx); err != nil {
+			t.Logf("Repetitive commit N %d", i+1)
+			t.Fatal(err)
+		}
+	}
+	committed, rolledBack, err := cleanup(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !committed {
+		t.Log("tx was not committed but we expected it to")
+		t.FailNow()
+	}
+	if rolledBack {
+		t.Log("tx was rolled back and we did not expect it to")
+		t.FailNow()
+	}
+
+	if fc.begin != 1 {
+		t.Logf("begin was called %d times in the underlying conn but we expected 1", fc.begin)
+		t.FailNow()
+	}
+
+	if fc.commit != 1 {
+		t.Logf("commit was called %d times in the underlying conn but we expected 1", fc.commit)
+		t.FailNow()
+	}
+
+	if fc.rollback != 0 {
+		t.Logf("rollback was called %d times in the underlying conn but we expected 0", fc.rollback)
+		t.FailNow()
+	}
+}