-
Notifications
You must be signed in to change notification settings - Fork 66
/
bulk_insert.go
143 lines (120 loc) · 4.25 KB
/
bulk_insert.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
package gormbulk
import (
"errors"
"fmt"
"reflect"
"strings"
"time"
"github.com/jinzhu/gorm"
)
// Insert multiple records at once
// [objects] Must be a slice of struct
// [chunkSize] Number of records to insert at once.
// Embedding a large number of variables at once will raise an error beyond the limit of prepared statement.
// Larger size will normally lead the better performance, but 2000 to 3000 is reasonable.
// [excludeColumns] Columns you want to exclude from insert. You can omit if there is no column you want to exclude.
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 {
return err
}
}
return nil
}
func insertObjSet(db *gorm.DB, objects []interface{}, excludeColumns ...string) error {
if len(objects) == 0 {
return nil
}
firstAttrs, err := extractMapValue(objects[0], excludeColumns)
if err != nil {
return err
}
attrSize := len(firstAttrs)
// Scope to eventually run SQL
mainScope := db.NewScope(objects[0])
// Store placeholders for embedding variables
placeholders := make([]string, 0, attrSize)
// Replace with database column name
dbColumns := make([]string, 0, attrSize)
for _, key := range sortedKeys(firstAttrs) {
dbColumns = append(dbColumns, mainScope.Quote(gorm.ToColumnName(key)))
}
for _, obj := range objects {
objAttrs, err := extractMapValue(obj, excludeColumns)
if err != nil {
return err
}
// If object sizes are different, SQL statement loses consistency
if len(objAttrs) != attrSize {
return errors.New("attribute sizes are inconsistent")
}
scope := db.NewScope(obj)
// Append variables
variables := make([]string, 0, attrSize)
for _, key := range sortedKeys(objAttrs) {
scope.AddToVars(objAttrs[key])
variables = append(variables, "?")
}
valueQuery := "(" + strings.Join(variables, ", ") + ")"
placeholders = append(placeholders, valueQuery)
// Also append variables to mainScope
mainScope.SQLVars = append(mainScope.SQLVars, scope.SQLVars...)
}
insertOption := ""
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")
}
insertOption = strVal
}
mainScope.Raw(fmt.Sprintf("INSERT INTO %s (%s) VALUES %s %s",
mainScope.QuotedTableName(),
strings.Join(dbColumns, ", "),
strings.Join(placeholders, ", "),
insertOption,
))
return db.Exec(mainScope.SQL, mainScope.SQLVars...).Error
}
// Obtain columns and values required for insert from interface
func extractMapValue(value interface{}, excludeColumns []string) (map[string]interface{}, error) {
rv := reflect.ValueOf(value)
if rv.Kind() == reflect.Ptr {
rv = rv.Elem()
value = rv.Interface()
}
if rv.Kind() != reflect.Struct {
return nil, errors.New("value must be kind of Struct")
}
var attrs = map[string]interface{}{}
for _, field := range (&gorm.Scope{Value: value}).Fields() {
// Exclude relational record because it's not directly contained in database columns
_, hasForeignKey := field.TagSettingsGet("FOREIGNKEY")
if !containString(excludeColumns, field.Struct.Name) && field.StructField.Relationship == nil && !hasForeignKey &&
!field.IsIgnored && !fieldIsAutoIncrement(field) && !fieldIsPrimaryAndBlank(field) {
if (field.Struct.Name == "CreatedAt" || field.Struct.Name == "UpdatedAt") && field.IsBlank {
attrs[field.DBName] = time.Now()
} else if field.StructField.HasDefaultValue && field.IsBlank {
// If default value presents and field is empty, assign a default value
if val, ok := field.TagSettingsGet("DEFAULT"); ok {
attrs[field.DBName] = val
} else {
attrs[field.DBName] = field.Field.Interface()
}
} else {
attrs[field.DBName] = field.Field.Interface()
}
}
}
return attrs, nil
}
func fieldIsAutoIncrement(field *gorm.Field) bool {
if value, ok := field.TagSettingsGet("AUTO_INCREMENT"); ok {
return strings.ToLower(value) != "false"
}
return false
}
func fieldIsPrimaryAndBlank(field *gorm.Field) bool {
return field.IsPrimaryKey && field.IsBlank
}