聊聊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