聊聊gorm的CreateInBatches

序

本文主要研究一下gorm的CreateInBatches

CreateInBatches

gorm.io/gorm@v1.20.11/finisher_api.go

// CreateInBatches insert the value in batches into database
func (db *DB) CreateInBatches(value interface{}, batchSize int) (tx *DB) {
    reflectValue := reflect.Indirect(reflect.ValueOf(value))

    switch reflectValue.Kind() {
    case reflect.Slice, reflect.Array:
        var rowsAffected int64
        tx = db.getInstance()
        tx.AddError(tx.Transaction(func(tx *DB) error {
            for i := 0; i < reflectValue.Len(); i += batchSize {
                ends := i + batchSize
                if ends > reflectValue.Len() {
                    ends = reflectValue.Len()
                }

                subtx := tx.getInstance()
                subtx.Statement.Dest = reflectValue.Slice(i, ends).Interface()
                subtx.callbacks.Create().Execute(subtx)
                if subtx.Error != nil {
                    return subtx.Error
                }
                rowsAffected += subtx.RowsAffected
            }
            return nil
        }))
        tx.RowsAffected = rowsAffected
    default:
        tx = db.getInstance()
        tx.Statement.Dest = value
        tx.callbacks.Create().Execute(tx)
    }
    return
}

CreateInBatches会根据batchSize来分配进行create，但是他们是在同一个事务的，其rowsAffected是每个批次的rowsAffected的累加

AddError

gorm.io/gorm@v1.20.11/gorm.go

// AddError add error to db
func (db *DB) AddError(err error) error {
    if db.Error == nil {
        db.Error = err
    } else if err != nil {
        db.Error = fmt.Errorf("%v; %w", db.Error, err)
    }
    return db.Error
}

tx.AddError方法在db.Error为nil时直接更新为err；不为nil时判断err是否为nil，不为nil时才更新为err

CommitOrRollbackTransaction

gorm.io/gorm@v1.20.11/callbacks/transaction.go

func CommitOrRollbackTransaction(db *gorm.DB) {
    if !db.Config.SkipDefaultTransaction {
        if _, ok := db.InstanceGet("gorm:started_transaction"); ok {
            if db.Error == nil {
                db.Commit()
            } else {
                db.Rollback()
            }
            db.Statement.ConnPool = db.ConnPool
        }
    }
}

CommitOrRollbackTransaction方法会判断db.Error，如果不为nil则执行db.Rollback()

实例

func createInBatchesDemo(db *gorm.DB) {
    entities := []DemoEntity{
        {
            Name: "coco",
        },
        {
            Name: "bear",
        },
    }
    result := db.Debug().CreateInBatches(&entities, 1)
    b, _ := json.Marshal(entities)
    log.Println("data:", string(b))
    log.Println("result.RowsAffected:", result.RowsAffected, "result.Error:", result.Error)
}

输出

2021/01/16 22:28:55 /demo.go:71
[0.384ms] [rows:1] INSERT INTO `demo_entities` (`created_at`,`updated_at`,`deleted_at`,`name`) VALUES ("2021-01-16 22:28:55.774","2021-01-16 22:28:55.774",NULL,"coco")

2021/01/16 22:28:55 /demo.go:71
[0.055ms] [rows:1] INSERT INTO `demo_entities` (`created_at`,`updated_at`,`deleted_at`,`name`) VALUES ("2021-01-16 22:28:55.774","2021-01-16 22:28:55.774",NULL,"bear")
2021/01/16 22:28:55 data: [{"ID":7,"CreatedAt":"2021-01-16T22:28:55.7743+08:00","UpdatedAt":"2021-01-16T22:28:55.7743+08:00","DeletedAt":null,"Name":"coco"},{"ID":8,"CreatedAt":"2021-01-16T22:28:55.774794+08:00","UpdatedAt":"2021-01-16T22:28:55.774794+08:00","DeletedAt":null,"Name":"bear"}]
2021/01/16 22:28:55 result.RowsAffected: 2 result.Error: <nil>

小结

gorm的CreateInBatches允许用户自定义batchSize，它会根据batchSize来分配进行create，但是他们是在同一个事务的，其rowsAffected是每个批次的rowsAffected的累加。

doc

• gorm

有疑问加站长微信联系（非本文作者）