Add CreateInBatches supports

This commit is contained in:
Jinzhu 2020-11-16 21:42:30 +08:00
parent 62be27d3ca
commit a8db54afd6
2 changed files with 49 additions and 0 deletions

View File

@ -21,6 +21,29 @@ func (db *DB) Create(value interface{}) (tx *DB) {
return return
} }
// 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:
tx = db.getInstance()
for i := 0; i < reflectValue.Len(); i += batchSize {
tx.AddError(tx.Transaction(func(tx *DB) error {
ends := i + batchSize
if ends > reflectValue.Len() {
ends = reflectValue.Len()
}
return tx.Create(reflectValue.Slice(i, ends).Interface()).Error
}))
}
default:
return db.Create(value)
}
return
}
// Save update value in database, if the value doesn't have primary key, will insert it // Save update value in database, if the value doesn't have primary key, will insert it
func (db *DB) Save(value interface{}) (tx *DB) { func (db *DB) Save(value interface{}) (tx *DB) {
tx = db.getInstance() tx = db.getInstance()

View File

@ -40,6 +40,32 @@ func TestCreate(t *testing.T) {
} }
} }
func TestCreateInBatches(t *testing.T) {
users := []User{
*GetUser("create_in_batches_1", Config{Account: true, Pets: 2, Toys: 3, Company: true, Manager: true, Team: 0, Languages: 1, Friends: 1}),
*GetUser("create_in_batches_2", Config{Account: false, Pets: 2, Toys: 4, Company: false, Manager: false, Team: 1, Languages: 3, Friends: 5}),
*GetUser("create_in_batches_3", Config{Account: true, Pets: 0, Toys: 3, Company: true, Manager: false, Team: 4, Languages: 0, Friends: 1}),
*GetUser("create_in_batches_4", Config{Account: true, Pets: 3, Toys: 0, Company: false, Manager: true, Team: 0, Languages: 3, Friends: 0}),
*GetUser("create_in_batches_5", Config{Account: false, Pets: 0, Toys: 3, Company: true, Manager: false, Team: 1, Languages: 3, Friends: 1}),
*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)
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 TestCreateFromMap(t *testing.T) { func TestCreateFromMap(t *testing.T) {
if err := DB.Model(&User{}).Create(map[string]interface{}{"Name": "create_from_map", "Age": 18}).Error; err != nil { if err := DB.Model(&User{}).Create(map[string]interface{}{"Name": "create_from_map", "Age": 18}).Error; err != nil {
t.Fatalf("failed to create data from map, got error: %v", err) t.Fatalf("failed to create data from map, got error: %v", err)