Modify interface and implementation of BulkInsertWithAssigningIDs

README.md

@@ -46,118 +46,68 @@ In the above pattern `Name` and `Email` fields are excluded.
 
 ## Example
 
-### BulkInsert
-
-```go
-package main
-
-import (
-	"fmt"
-	"log"
-	"time"
-
-	"github.com/jinzhu/gorm"
-	_ "github.com/jinzhu/gorm/dialects/mysql"
-	gormbulk "github.com/t-tiger/gorm-bulk-insert/v2"
-)
-
-type fakeTable struct {
-	ID        int `gorm:"AUTO_INCREMENT"`
-	Name      string
-	Email     string
-	CreatedAt time.Time
-	UpdatedAt time.Time
-}
-
-func main() {
-	db, err := gorm.Open("mysql", "mydb")
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	var insertRecords []interface{}
-	for i := 0; i < 10; i++ {
-		insertRecords = append(insertRecords,
-			fakeTable{
-				Name:  fmt.Sprintf("name%d", i),
-				Email: fmt.Sprintf("test%[email protected]", i),
-				// you don't need to set CreatedAt, UpdatedAt
-			},
-		)
-	}
-
-	err = gormbulk.BulkInsert(db, insertRecords, 3000)
-	if err != nil {
-		// do something
-	}
-
-	// columns you want to exclude from Insert, specify as an argument
-	err = gormbulk.BulkInsert(db, insertRecords, 3000, "Email")
-	if err != nil {
-		// do something
-	}
-}
-```
-
-### BulkInsertWithAssigningIDs
-
 ```go
 package main
 
 import (
-	"fmt"
-	"time"
+  "fmt"
+  "log"
+  "time"
 
-	"github.com/jinzhu/gorm"
-	_ "github.com/jinzhu/gorm/dialects/postgres"
-	gormbulk "github.com/t-tiger/gorm-bulk-insert/v2"
+  "github.com/jinzhu/gorm"
+  _ "github.com/jinzhu/gorm/dialects/mysql"
+  gormbulk "github.com/t-tiger/gorm-bulk-insert/v2"
 )
 
 type fakeTable struct {
-	IdPK      uint      `gorm:"primary_key"`
-	CreatedAt time.Time `gorm:"default:now()"`
-	Data      string
+  ID        int `gorm:"AUTO_INCREMENT"`
+  Name      string
+  Email     string
+  CreatedAt time.Time
+  UpdatedAt time.Time
 }
 
 func main() {
-	db, err := gorm.Open("postgres", "host=localhost port=5432 user=cloudwalker dbname=cloudwalker password=cloudwalker sslmode=disable")
-	if err != nil {
-		panic(err)
-	}
-	defer db.Close()
-	db.SingularTable(true)
-
-	if err := db.AutoMigrate(fakeTable{}).Error; err != nil {
-		panic(err)
-	}
-
-	models := []interface{}{
-		fakeTable{Data: "aaa"},
-		fakeTable{Data: "bbb"},
-		fakeTable{Data: "ccc"},
-	}
-
-	// if you want to scan * back
-	var returnModel []fakeTable
-	if err := gormbulk.BulkInsertWithAssigningIDs(
-		db.Set("gorm:insert_option", "returning *"), &returnModel, models, 1000); err != nil {
-		panic(err)
-	}
-	fmt.Printf("success to insert with returning: %+v\n", returnModel)
-	// success to insert with returning: [
-	// {IdPK:1 CreatedAt:2021-10-31 16:21:48.019947 +0000 UTC Data:aaa} 
-	// {IdPK:2 CreatedAt:2021-10-31 16:21:48.019959 +0000 UTC Data:bbb} 
-	// {IdPK:3 CreatedAt:2021-10-31 16:21:48.019965 +0000 UTC Data:ccc}
-	// ]
-
-	// if you want to scan primary key
-	var returnId []uint
-	if err := gormbulk.BulkInsertWithAssigningIDs(
-		db.Set("gorm:insert_option", "returning id"), &returnId, models, 1000); err != nil {
-		panic(err)
-	}
-	fmt.Printf("success to insert with returning: %+v\n", returnId)
-	// `success to insert with returning: [4 5 6]`
+  db, err := gorm.Open("mysql", "mydb")
+  if err != nil {
+    log.Fatal(err)
+  }
+
+  var insertRecords []interface{}
+  for i := 0; i < 10; i++ {
+    insertRecords = append(insertRecords,
+      fakeTable{
+        Name:  fmt.Sprintf("name%d", i),
+        Email: fmt.Sprintf("test%[email protected]", i),
+        // you don't need to set CreatedAt, UpdatedAt
+      },
+    )
+  }
+
+  err = gormbulk.BulkInsert(db, insertRecords, 3000)
+  if err != nil {
+    // do something
+  }
+
+  // Columns you want to exclude from Insert, specify as an argument
+  err = gormbulk.BulkInsert(db, insertRecords, 3000, "Email")
+  if err != nil {
+    // do something
+  }
+
+  // Fetch returning values
+  dbForReturning := db.Set("gorm:insert_option", "RETURNING id, name, created_at")
+  var returned []struct {
+    ID        int
+    Name      string
+    CreatedAt time.Time
+  }
+  err = gormbulk.BulkInsertWithReturningValues(dbForReturning, insertRecords, &returned, 3000)
+  if err != nil {
+    // do something
+  }
+  // Values of `returned` will be as follows
+  // {{ID: 1, Name: "name0", CreatedAt: 2021-10-31 16:21:48.019947 +0000 UTC}, ...}
 }
 ```
 

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.
 //
+// [returnedVals] must be a point to a slice of struct. Values returned from `RETURNING` clause will be assigned.
+//
 // [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