Add CreateBatchSize option

This commit is contained in:
Jinzhu 2020-12-02 14:59:50 +08:00
parent 41e52f343a
commit 0c12a4c360
4 changed files with 63 additions and 9 deletions

View File

@ -15,6 +15,10 @@ import (
// Create insert the value into database
func (db *DB) Create(value interface{}) (tx *DB) {
if db.CreateBatchSize > 0 {
return db.CreateInBatches(value, db.CreateBatchSize)
}
tx = db.getInstance()
tx.Statement.Dest = value
tx.callbacks.Create().Execute(tx)
@ -27,19 +31,30 @@ func (db *DB) CreateInBatches(value interface{}, batchSize int) (tx *DB) {
switch reflectValue.Kind() {
case reflect.Slice, reflect.Array:
var rowsAffected int64
tx = db.getInstance()
for i := 0; i < reflectValue.Len(); i += batchSize {
tx.AddError(tx.Transaction(func(tx *DB) error {
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()
}
return tx.Create(reflectValue.Slice(i, ends).Interface()).Error
}))
}
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:
return db.Create(value)
tx = db.getInstance()
tx.Statement.Dest = value
tx.callbacks.Create().Execute(tx)
}
return
}

View File

@ -38,6 +38,8 @@ type Config struct {
AllowGlobalUpdate bool
// QueryFields executes the SQL query with all fields of the table
QueryFields bool
// CreateBatchSize default create batch size
CreateBatchSize int
// ClauseBuilders clause builder
ClauseBuilders map[string]clause.ClauseBuilder
@ -74,6 +76,7 @@ type Session struct {
Context context.Context
Logger logger.Interface
NowFunc func() time.Time
CreateBatchSize int
}
// Open initialize db session based on dialector
@ -161,6 +164,10 @@ func (db *DB) Session(config *Session) *DB {
}
)
if config.CreateBatchSize > 0 {
tx.Config.CreateBatchSize = config.CreateBatchSize
}
if config.SkipDefaultTransaction {
tx.Config.SkipDefaultTransaction = true
}

View File

@ -50,7 +50,39 @@ func TestCreateInBatches(t *testing.T) {
*GetUser("create_in_batches_6", Config{Account: true, Pets: 4, Toys: 3, Company: false, Manager: true, Team: 1, Languages: 3, Friends: 0}),
}
DB.CreateInBatches(&users, 2)
result := DB.CreateInBatches(&users, 2)
if result.RowsAffected != int64(len(users)) {
t.Errorf("affected rows should be %v, but got %v", len(users), result.RowsAffected)
}
for _, user := range users {
if user.ID == 0 {
t.Fatalf("failed to fill user's ID, got %v", user.ID)
} else {
var newUser User
if err := DB.Where("id = ?", user.ID).Preload(clause.Associations).First(&newUser).Error; err != nil {
t.Fatalf("errors happened when query: %v", err)
} else {
CheckUser(t, newUser, user)
}
}
}
}
func TestCreateInBatchesWithDefaultSize(t *testing.T) {
users := []User{
*GetUser("create_with_default_batch_size_1", Config{Account: true, Pets: 2, Toys: 3, Company: true, Manager: true, Team: 0, Languages: 1, Friends: 1}),
*GetUser("create_with_default_batch_sizs_2", Config{Account: false, Pets: 2, Toys: 4, Company: false, Manager: false, Team: 1, Languages: 3, Friends: 5}),
*GetUser("create_with_default_batch_sizs_3", Config{Account: true, Pets: 0, Toys: 3, Company: true, Manager: false, Team: 4, Languages: 0, Friends: 1}),
*GetUser("create_with_default_batch_sizs_4", Config{Account: true, Pets: 3, Toys: 0, Company: false, Manager: true, Team: 0, Languages: 3, Friends: 0}),
*GetUser("create_with_default_batch_sizs_5", Config{Account: false, Pets: 0, Toys: 3, Company: true, Manager: false, Team: 1, Languages: 3, Friends: 1}),
*GetUser("create_with_default_batch_sizs_6", Config{Account: true, Pets: 4, Toys: 3, Company: false, Manager: true, Team: 1, Languages: 3, Friends: 0}),
}
result := DB.Session(&gorm.Session{CreateBatchSize: 2}).Create(&users)
if result.RowsAffected != int64(len(users)) {
t.Errorf("affected rows should be %v, but got %v", len(users), result.RowsAffected)
}
for _, user := range users {
if user.ID == 0 {

View File

@ -9,9 +9,9 @@ require (
github.com/stretchr/testify v1.5.1
gorm.io/driver/mysql v1.0.3
gorm.io/driver/postgres v1.0.5
gorm.io/driver/sqlite v1.1.3
gorm.io/driver/sqlite v1.1.4
gorm.io/driver/sqlserver v1.0.5
gorm.io/gorm v1.20.5
gorm.io/gorm v1.20.7
)
replace gorm.io/gorm => ../