github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/sqlx/tx_task.go (about)

     1  package sqlx
     2  
     3  import (
     4  	"fmt"
     5  	"runtime/debug"
     6  
     7  	"github.com/pkg/errors"
     8  
     9  	"github.com/machinefi/w3bstream/pkg/depends/kit/logr"
    10  )
    11  
    12  type Task func(db DBExecutor) error
    13  
    14  func (task Task) Run(db DBExecutor) (err error) {
    15  	defer func() {
    16  		if e := recover(); e != nil {
    17  			err = fmt.Errorf("panic: %s; calltrace:%s", fmt.Sprint(e), string(debug.Stack()))
    18  		}
    19  	}()
    20  	return task(db)
    21  }
    22  
    23  func NewTasks(db DBExecutor) *Tasks {
    24  	return &Tasks{
    25  		db: db,
    26  	}
    27  }
    28  
    29  type Tasks struct {
    30  	db    DBExecutor
    31  	tasks []Task
    32  }
    33  
    34  func (tasks Tasks) With(task ...Task) *Tasks {
    35  	tasks.tasks = append(tasks.tasks, task...)
    36  	return &tasks
    37  }
    38  
    39  func (tasks *Tasks) Do() (err error) {
    40  	if len(tasks.tasks) == 0 {
    41  		return nil
    42  	}
    43  
    44  	db := tasks.db
    45  
    46  	l := logr.FromContext(db.Context())
    47  
    48  	if maybeTx, ok := db.(TxExecutor); ok {
    49  		inTxScope := false
    50  
    51  		if !maybeTx.IsTx() {
    52  			db, err = maybeTx.Begin()
    53  			if err != nil {
    54  				return err
    55  			}
    56  			maybeTx = db.(TxExecutor)
    57  			inTxScope = true
    58  		}
    59  
    60  		for _, task := range tasks.tasks {
    61  			if runErr := task.Run(db); runErr != nil {
    62  				if inTxScope {
    63  					// err will bubble up，just handle and rollback in outermost layer
    64  					l.Error(errors.Wrap(err, "SQL FAILED"))
    65  					if rollBackErr := maybeTx.Rollback(); rollBackErr != nil {
    66  						l.Warn(errors.Wrap(rollBackErr, "ROLLBACK FAILED"))
    67  						err = rollBackErr
    68  						return
    69  					}
    70  				}
    71  				return runErr
    72  			}
    73  		}
    74  
    75  		if inTxScope {
    76  			if commitErr := maybeTx.Commit(); commitErr != nil {
    77  				l.Warn(errors.Wrap(commitErr, "TRANSACTION COMMIT FAILED"))
    78  				return commitErr
    79  			}
    80  		}
    81  
    82  	}
    83  	return nil
    84  }