cleaned up transaction logic

0chain · Aug 4, 2023 · ab5dc52 · ab5dc52
1 parent 3138346
commit ab5dc52
Show file tree

Hide file tree

Showing 13 changed files with 122 additions and 217 deletions.
diff --git a/code/go/0chain.net/blobbercore/allocation/dao.go b/code/go/0chain.net/blobbercore/allocation/dao.go
@@ -3,7 +3,6 @@ package allocation
 import (
 	"context"
 
-	"github.com/0chain/blobber/code/go/0chain.net/blobbercore/datastore"
 	"github.com/0chain/blobber/code/go/0chain.net/core/common"
 	"github.com/0chain/errors"
 	"github.com/0chain/gosdk/constants"
@@ -12,16 +11,11 @@ import (
 
 // GetOrCreate, get allocation if it exists in db. if not, try to sync it from blockchain, and insert it in db.
 func GetOrCreate(ctx context.Context, allocationId string) (*Allocation, error) {
-
-	db := datastore.GetStore().CreateTransaction(ctx)
-	tx := datastore.GetStore().GetTransaction(ctx)
-
 	if len(allocationId) == 0 {
 		return nil, errors.Throw(constants.ErrInvalidParameter, "tx")
 	}
 
-	alloc, err := Repo.GetById(db, allocationId)
-	tx.Rollback()
+	alloc, err := Repo.GetById(ctx, allocationId)
 
 	if err == nil {
 		return alloc, nil

diff --git a/code/go/0chain.net/blobbercore/allocation/workers.go b/code/go/0chain.net/blobbercore/allocation/workers.go
@@ -13,8 +13,6 @@ import (
 	"github.com/0chain/blobber/code/go/0chain.net/core/logging"
 	"github.com/0chain/blobber/code/go/0chain.net/core/transaction"
 
-	"gorm.io/gorm"
-
 	"go.uber.org/zap"
 )
 
@@ -169,14 +167,6 @@ func requestAllocation(allocID string) (sa *transaction.StorageAllocation, err e
 	return
 }
 
-func commit(tx *gorm.DB, err *error) {
-	if (*err) != nil {
-		tx.Rollback()
-		return
-	}
-	(*err) = tx.Commit().Error
-}
-
 func updateAllocationInDB(ctx context.Context, a *Allocation, sa *transaction.StorageAllocation) (ua *Allocation, err error) {
 	var tx = datastore.GetStore().GetTransaction(ctx)
 	var changed bool = a.Tx != sa.Tx

diff --git a/code/go/0chain.net/blobbercore/automigration/automigration.go b/code/go/0chain.net/blobbercore/automigration/automigration.go
diff --git a/code/go/0chain.net/blobbercore/datastore/mocket.go b/code/go/0chain.net/blobbercore/datastore/mocket.go
@@ -99,6 +99,22 @@ func (store *Mocket) WithNewTransaction(f func(ctx context.Context) error) error
 	return nil
 }
 
+func (store *Mocket) WithTransaction(ctx context.Context, f func(ctx context.Context) error) error {
+	tx := store.GetTransaction(ctx)
+	if tx == nil {
+		ctx = store.CreateTransaction(ctx)
+		tx = store.GetTransaction(ctx)
+	}
+
+	err := f(ctx)
+	if err != nil {
+		tx.Rollback()
+		return err
+	}
+	tx.Commit()
+	return nil
+}
+
 func (store *Mocket) GetDB() *gorm.DB {
 	return store.db
 }

diff --git a/code/go/0chain.net/blobbercore/datastore/postgres.go b/code/go/0chain.net/blobbercore/datastore/postgres.go
@@ -116,6 +116,21 @@ func (store *postgresStore) WithNewTransaction(f func(ctx context.Context) error
 	tx.Commit()
 	return nil
 }
+func (store *postgresStore) WithTransaction(ctx context.Context, f func(ctx context.Context) error) error {
+	tx := store.GetTransaction(ctx)
+	if tx == nil {
+		ctx = store.CreateTransaction(ctx)
+		tx = store.GetTransaction(ctx)
+	}
+
+	err := f(ctx)
+	if err != nil {
+		tx.Rollback()
+		return err
+	}
+	tx.Commit()
+	return nil
+}
 
 func (store *postgresStore) GetDB() *gorm.DB {
 	return store.db

diff --git a/code/go/0chain.net/blobbercore/datastore/sqlmock.go b/code/go/0chain.net/blobbercore/datastore/sqlmock.go
@@ -94,6 +94,22 @@ func (store *Sqlmock) WithNewTransaction(f func(ctx context.Context) error) erro
 	return nil
 }
 
+func (store *Sqlmock) WithTransaction(ctx context.Context, f func(ctx context.Context) error) error {
+	tx := store.GetTransaction(ctx)
+	if tx == nil {
+		ctx = store.CreateTransaction(ctx)
+		tx = store.GetTransaction(ctx)
+	}
+
+	err := f(ctx)
+	if err != nil {
+		tx.Rollback()
+		return err
+	}
+	tx.Commit()
+	return nil
+}
+
 func (store *Sqlmock) GetDB() *gorm.DB {
 	return store.db
 }

diff --git a/code/go/0chain.net/blobbercore/datastore/store.go b/code/go/0chain.net/blobbercore/datastore/store.go
@@ -32,6 +32,7 @@ type Store interface {
 	// GetTransaction get transaction from context
 	GetTransaction(ctx context.Context) *EnhancedDB
 	WithNewTransaction(f func(ctx context.Context) error) error
+	WithTransaction(ctx context.Context, f func(ctx context.Context) error) error
 	// Get db connection with user that creates roles and databases. Its dialactor does not contain database name
 	GetPgDB() (*gorm.DB, error)
 	Open() error

diff --git a/code/go/0chain.net/blobbercore/handler/context.go b/code/go/0chain.net/blobbercore/handler/context.go
@@ -130,8 +130,8 @@ type ErrorResponse struct {
 	Error string
 }
 
-// WithHandler process handler to respond request
-func WithHandler(handler func(ctx *Context) (interface{}, error)) func(w http.ResponseWriter, r *http.Request) {
+// WithTxHandler process handler to respond request
+func WithTxHandler(handler func(ctx *Context) (interface{}, error)) func(w http.ResponseWriter, r *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Access-Control-Allow-Origin", "*") // CORS for all.
 		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
@@ -143,23 +143,39 @@ func WithHandler(handler func(ctx *Context) (interface{}, error)) func(w http.Re
 		}
 
 		common.TryParseForm(r)
-
 		w.Header().Set("Content-Type", "application/json")
 
-		ctx, err := WithVerify(r)
-		statusCode := ctx.StatusCode
+		statusCode := 0
+		var result interface{}
+		err := datastore.GetStore().WithNewTransaction(func(c context.Context) error {
+			ctx := &Context{
+				Context: c,
+				Request: r,
+				Store:   datastore.GetStore(),
+			}
 
-		if err != nil {
-			if statusCode == 0 {
-				statusCode = http.StatusInternalServerError
+			ctx.Vars = mux.Vars(r)
+			if ctx.Vars == nil {
+				ctx.Vars = make(map[string]string)
 			}
 
-			http.Error(w, err.Error(), statusCode)
-			return
-		}
+			ctx.ClientID = r.Header.Get(common.ClientHeader)
+			ctx.ClientKey = r.Header.Get(common.ClientKeyHeader)
+			ctx.AllocationId = r.Header.Get(common.AllocationIdHeader)
+			ctx.Signature = r.Header.Get(common.ClientSignatureHeader)
 
-		result, err := handler(ctx)
-		statusCode = ctx.StatusCode
+			ctx, err := WithVerify(ctx, r)
+			statusCode = ctx.StatusCode
+
+			if err != nil {
+				return err
+			}
+
+			result, err = handler(ctx)
+			statusCode = ctx.StatusCode
+
+			return nil
+		})
 
 		if err != nil {
 			if statusCode == 0 {
@@ -183,23 +199,8 @@ func WithHandler(handler func(ctx *Context) (interface{}, error)) func(w http.Re
 }
 
 // WithVerify verify allocation and signature
-func WithVerify(r *http.Request) (*Context, error) {
-
-	ctx := &Context{
-		Context: context.TODO(),
-		Request: r,
-		Store:   datastore.GetStore(),
-	}
-
-	ctx.Vars = mux.Vars(r)
-	if ctx.Vars == nil {
-		ctx.Vars = make(map[string]string)
-	}
+func WithVerify(ctx *Context, r *http.Request) (*Context, error) {
 
-	ctx.ClientID = r.Header.Get(common.ClientHeader)
-	ctx.ClientKey = r.Header.Get(common.ClientKeyHeader)
-	ctx.AllocationId = r.Header.Get(common.AllocationIdHeader)
-	ctx.Signature = r.Header.Get(common.ClientSignatureHeader)
 	allocationTx := ctx.Vars["allocation"]
 
 	if len(ctx.AllocationId) > 0 {