github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/sqlx/task.go (about)

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