Modify interface and implementation of BulkInsertWithAssigningIDs

bulk_insert.go

@@ -25,79 +25,58 @@ import (
 func BulkInsert(db *gorm.DB, objects []interface{}, chunkSize int, excludeColumns ...string) error {
 	// Split records with specified size not to exceed Database parameter limit
 	for _, objSet := range splitObjects(objects, chunkSize) {
-		if err := insertObjSet(db, objSet, excludeColumns...); err != nil {
+		_, err := insertObjSet(db, objSet, excludeColumns...)
+		if err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-// BulkInsertWithAssigningIDs executes the query to insert multiple records at once.
-// it will scan the result of `returning id` or `returning *` to [returnedValue] after every insert.
-// it's necessary to set "gorm:insert_option"="returning id" in *gorm.DB
-//
-// [returnedValue] slice of primary_key or model, must be a *[]uint(for integer), *[]string(for uuid), *[]struct(for `returning *`)
+// BulkInsertWithReturningValues executes the query to insert multiple records at once.
+// This will scan the returned value into `dstValues`.
+// It's necessary to set "gorm:insert_option" to execute "returning" query.
 //
 // [objects] must be a slice of struct.
 //
+// [dstValues] must be a point to a slice of struct. Struct properties must correspond to returning results.
+//
 // [chunkSize] is a number of variables embedded in query. To prevent the error which occurs embedding a large number of variables at once
 // and exceeds the limit of prepared statement. Larger size normally leads to better performance, in most cases 2000 to 3000 is reasonable.
 //
 // [excludeColumns] is column names to exclude from insert.
-func BulkInsertWithAssigningIDs(db *gorm.DB, returnedValue interface{}, objects []interface{}, chunkSize int, excludeColumns ...string) error {
-	typ := reflect.TypeOf(returnedValue)
-	if typ.Kind() != reflect.Ptr || typ.Elem().Kind() != reflect.Slice {
-		return errors.New("returningId must be a slice ptr")
+func BulkInsertWithReturningValues(db *gorm.DB, objects []interface{}, returnedVals interface{}, chunkSize int, excludeColumns ...string) error {
+	typ := reflect.TypeOf(returnedVals)
+	if typ.Kind() != reflect.Ptr || typ.Elem().Kind() != reflect.Slice || typ.Elem().Elem().Kind() != reflect.Struct {
+		return errors.New("returnedVals must be a pointer to a slice of struct")
 	}
 
-	allIds := reflect.Indirect(reflect.ValueOf(returnedValue))
-	typ = allIds.Type()
-
-	// Deference value of slice
-	valueTyp := typ.Elem()
-	for valueTyp.Kind() == reflect.Ptr {
-		valueTyp = valueTyp.Elem()
-	}
+	refDst := reflect.Indirect(reflect.ValueOf(returnedVals))
+	typ = refDst.Type()
 
 	// Split records with specified size not to exceed Database parameter limit
 	for _, objSet := range splitObjects(objects, chunkSize) {
-		returnValueSlice := reflect.New(typ)
-		var scanReturningId func(*gorm.DB) error
-		switch valueTyp.Kind() {
-		case reflect.Struct:
-			// If user want to scan `returning *` with returnedValue=[]struct{...}
-			scanReturningId = func(db *gorm.DB) error {
-				return db.Scan(returnValueSlice.Interface()).Error
-			}
-		default:
-			// If user want to scan primary key `returning pk` with returnedValue=[]struct{...}
-			pk := db.NewScope(objects[0]).PrimaryKey()
-			scanReturningId = func(db *gorm.DB) error {
-				return db.Pluck(pk, returnValueSlice.Interface()).Error
-			}
+		db, err := insertObjSet(db, objSet, excludeColumns...)
+		if err != nil {
+			return err
 		}
-
-		if err := insertObjSetWithCallback(db, objSet, scanReturningId, excludeColumns...); err != nil {
+		scanned := reflect.New(typ)
+		if err := db.Scan(scanned.Interface()).Error; err != nil {
 			return err
 		}
-
-		allIds.Set(reflect.AppendSlice(allIds, returnValueSlice.Elem()))
+		refDst.Set(reflect.AppendSlice(refDst, scanned.Elem()))
 	}
 	return nil
 }
 
-func insertObjSet(db *gorm.DB, objects []interface{}, excludeColumns ...string) error {
-	return insertObjSetWithCallback(db, objects, nil, excludeColumns...)
-}
-
-func insertObjSetWithCallback(db *gorm.DB, objects []interface{}, postInsert func(*gorm.DB) error, excludeColumns ...string) error {
+func insertObjSet(db *gorm.DB, objects []interface{}, excludeColumns ...string) (*gorm.DB, error) {
 	if len(objects) == 0 {
-		return nil
+		return db, nil
 	}
 
 	firstAttrs, err := extractMapValue(objects[0], excludeColumns)
 	if err != nil {
-		return err
+		return db, err
 	}
 
 	attrSize := len(firstAttrs)
@@ -116,12 +95,12 @@ func insertObjSetWithCallback(db *gorm.DB, objects []interface{}, postInsert fun
 	for _, obj := range objects {
 		objAttrs, err := extractMapValue(obj, excludeColumns)
 		if err != nil {
-			return err
+			return db, err
 		}
 
 		// If object sizes are different, SQL statement loses consistency
 		if len(objAttrs) != attrSize {
-			return errors.New("attribute sizes are inconsistent")
+			return db, errors.New("attribute sizes are inconsistent")
 		}
 
 		scope := db.NewScope(obj)
@@ -144,7 +123,7 @@ func insertObjSetWithCallback(db *gorm.DB, objects []interface{}, postInsert fun
 	if val, ok := db.Get("gorm:insert_option"); ok {
 		strVal, ok := val.(string)
 		if !ok {
-			return errors.New("gorm:insert_option should be a string")
+			return db, errors.New("gorm:insert_option should be a string")
 		}
 		insertOption = strVal
 	}
@@ -157,18 +136,10 @@ func insertObjSetWithCallback(db *gorm.DB, objects []interface{}, postInsert fun
 	))
 
 	db = db.Raw(mainScope.SQL, mainScope.SQLVars...)
-
 	if err := db.Error; err != nil {
-		return err
-	}
-
-	if postInsert != nil {
-		if err := postInsert(db); err != nil {
-			return err
-		}
+		return db, err
 	}
-
-	return nil
+	return db, nil
 }
 
 // Obtain columns and values required for insert from interface

bulk_insert_test.go

@@ -26,6 +26,55 @@ type fakeTable struct {
 	UpdatedAt time.Time
 }
 
+func TestBulkInsertWithReturningValues(t *testing.T) {
+	type Table struct {
+		ID            uint `gorm:"primary_key;auto_increment"`
+		RegularColumn string
+		Custom        string `gorm:"column:ThisIsCamelCase"`
+	}
+
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+
+	defer db.Close()
+
+	gdb, err := gorm.Open("mysql", db)
+	require.NoError(t, err)
+
+	mock.ExpectQuery(
+		"INSERT INTO `tables` \\(`ThisIsCamelCase`, `regular_column`\\)",
+	).WithArgs(
+		"first custom", "first regular",
+		"second custom", "second regular",
+	).WillReturnRows(
+		sqlmock.NewRows([]string{"id", "ThisIsCamelCase", "regular_column"}).
+			AddRow(1, "first custom", "first regular").
+			AddRow(2, "second custom", "second regular"),
+	)
+
+	var returnedVals []Table
+	obj := []interface{}{
+		Table{
+			RegularColumn: "first regular",
+			Custom:        "first custom",
+		},
+		Table{
+			RegularColumn: "second regular",
+			Custom:        "second custom",
+		},
+	}
+
+	gdb = gdb.Set("gorm_insert_option", "RETURNING id, ThisIsCamelCase, regular_column")
+	err = BulkInsertWithReturningValues(gdb, obj, &returnedVals, 1000)
+	require.NoError(t, err)
+
+	expected := []Table{
+		{ID: 1, RegularColumn: "first regular", Custom: "first custom"},
+		{ID: 2, RegularColumn: "second regular", Custom: "second custom"},
+	}
+	assert.Equal(t, expected, returnedVals)
+}
+
 func Test_extractMapValue(t *testing.T) {
 	collectKeys := func(val map[string]interface{}) []string {
 		keys := make([]string, 0, len(val))
@@ -99,7 +148,7 @@ func Test_insertObject(t *testing.T) {
 		sqlmock.NewResult(1, 1),
 	)
 
-	err = insertObjSet(gdb, []interface{}{
+	_, err = insertObjSet(gdb, []interface{}{
 		Table{
 			RegularColumn: "first regular",
 			Custom:        "first custom",
@@ -113,59 +162,6 @@ func Test_insertObject(t *testing.T) {
 	require.NoError(t, err)
 }
 
-func Test_insertObjSetWithCallback(t *testing.T) {
-	type Table struct {
-		ID            uint `gorm:"primary_key;auto_increment"`
-		RegularColumn string
-		Custom        string `gorm:"column:ThisIsCamelCase"`
-	}
-
-	db, mock, err := sqlmock.New()
-	require.NoError(t, err)
-
-	defer db.Close()
-
-	gdb, err := gorm.Open("mysql", db)
-	require.NoError(t, err)
-
-	mock.ExpectQuery(
-		"INSERT INTO `tables` \\(`ThisIsCamelCase`, `regular_column`\\)",
-	).WithArgs(
-		"first custom", "first regular",
-		"second custom", "second regular",
-	).WillReturnRows(
-		sqlmock.NewRows([]string{"id"}).AddRow(1).AddRow(2),
-	)
-
-	returningIdScope := func(db *gorm.DB) *gorm.DB {
-		return db.Set("gorm:insert_option", "returning id")
-	}
-
-	err = insertObjSetWithCallback(gdb.Scopes(returningIdScope), []interface{}{
-		Table{
-			RegularColumn: "first regular",
-			Custom:        "first custom",
-		},
-		Table{
-			RegularColumn: "second regular",
-			Custom:        "second custom",
-		},
-	}, func(db *gorm.DB) error {
-		var ids []uint
-		if err := db.Pluck("id", &ids).Error; err != nil {
-			return err
-		}
-		require.Len(t, ids, 2, "must return 2 ids")
-		return nil
-	})
-
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	require.NoError(t, err)
-}
-
 func Test_fieldIsAutoIncrement(t *testing.T) {
 	type explicitSetTable struct {
 		ID int `gorm:"column:id;auto_increment"`
@@ -195,27 +191,3 @@ func Test_fieldIsAutoIncrement(t *testing.T) {
 		}
 	}
 }
-
-func Test_fieldIsPrimaryAndBlank(t *testing.T) {
-	type notPrimaryTable struct {
-		Dummy int
-	}
-	type primaryKeyTable struct {
-		ID int `gorm:"column:id;primary_key"`
-	}
-
-	cases := []struct {
-		Value    interface{}
-		Expected bool
-	}{
-		{notPrimaryTable{Dummy: 0}, false},
-		{notPrimaryTable{Dummy: 1}, false},
-		{primaryKeyTable{ID: 0}, true},
-		{primaryKeyTable{ID: 1}, false},
-	}
-	for _, c := range cases {
-		for _, field := range (&gorm.Scope{Value: c.Value}).Fields() {
-			assert.Equal(t, fieldIsPrimaryAndBlank(field), c.Expected)
-		}
-	}
-}