mirror of https://github.com/go-gorm/gorm.git
317 lines
9.0 KiB
Go
317 lines
9.0 KiB
Go
package callbacks
|
|
|
|
import (
|
|
"reflect"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
"gorm.io/gorm/schema"
|
|
"gorm.io/gorm/utils"
|
|
)
|
|
|
|
func SaveBeforeAssociations(db *gorm.DB) {
|
|
if db.Error == nil && db.Statement.Schema != nil {
|
|
selectColumns, restricted := SelectAndOmitColumns(db.Statement, true, false)
|
|
|
|
// Save Belongs To associations
|
|
for _, rel := range db.Statement.Schema.Relationships.BelongsTo {
|
|
if !saveAssociationCheck(db, rel, selectColumns, restricted) {
|
|
continue
|
|
}
|
|
|
|
setupReferences := func(obj reflect.Value, elem reflect.Value) {
|
|
for _, ref := range rel.References {
|
|
if !ref.OwnPrimaryKey {
|
|
pv, _ := ref.PrimaryKey.ValueOf(elem)
|
|
ref.ForeignKey.Set(obj, pv)
|
|
|
|
if dest, ok := db.Statement.Dest.(map[string]interface{}); ok {
|
|
dest[ref.ForeignKey.DBName] = pv
|
|
if _, ok := dest[rel.Name]; ok {
|
|
dest[rel.Name] = elem.Interface()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
switch db.Statement.ReflectValue.Kind() {
|
|
case reflect.Slice, reflect.Array:
|
|
var (
|
|
objs []reflect.Value
|
|
fieldType = rel.Field.FieldType
|
|
isPtr = fieldType.Kind() == reflect.Ptr
|
|
)
|
|
|
|
if !isPtr {
|
|
fieldType = reflect.PtrTo(fieldType)
|
|
}
|
|
|
|
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
|
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
|
obj := db.Statement.ReflectValue.Index(i)
|
|
if _, zero := rel.Field.ValueOf(obj); !zero { // check belongs to relation value
|
|
rv := rel.Field.ReflectValueOf(obj) // relation reflect value
|
|
objs = append(objs, obj)
|
|
if isPtr {
|
|
elems = reflect.Append(elems, rv)
|
|
} else {
|
|
elems = reflect.Append(elems, rv.Addr())
|
|
}
|
|
}
|
|
}
|
|
|
|
if elems.Len() > 0 {
|
|
if db.AddError(db.Session(&gorm.Session{}).Clauses(clause.OnConflict{
|
|
DoNothing: true,
|
|
}).Create(elems.Interface()).Error) == nil {
|
|
for i := 0; i < elems.Len(); i++ {
|
|
setupReferences(objs[i], elems.Index(i))
|
|
}
|
|
}
|
|
}
|
|
case reflect.Struct:
|
|
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
|
rv := rel.Field.ReflectValueOf(db.Statement.ReflectValue) // relation reflect value
|
|
if rv.Kind() != reflect.Ptr {
|
|
rv = rv.Addr()
|
|
}
|
|
|
|
if db.AddError(db.Session(&gorm.Session{}).Clauses(clause.OnConflict{
|
|
DoNothing: true,
|
|
}).Create(rv.Interface()).Error) == nil {
|
|
setupReferences(db.Statement.ReflectValue, rv)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func SaveAfterAssociations(db *gorm.DB) {
|
|
if db.Error == nil && db.Statement.Schema != nil {
|
|
selectColumns, restricted := SelectAndOmitColumns(db.Statement, true, false)
|
|
|
|
// Save Has One associations
|
|
for _, rel := range db.Statement.Schema.Relationships.HasOne {
|
|
if !saveAssociationCheck(db, rel, selectColumns, restricted) {
|
|
continue
|
|
}
|
|
|
|
switch db.Statement.ReflectValue.Kind() {
|
|
case reflect.Slice, reflect.Array:
|
|
var (
|
|
fieldType = rel.Field.FieldType
|
|
isPtr = fieldType.Kind() == reflect.Ptr
|
|
)
|
|
|
|
if !isPtr {
|
|
fieldType = reflect.PtrTo(fieldType)
|
|
}
|
|
|
|
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
|
|
|
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
|
obj := db.Statement.ReflectValue.Index(i)
|
|
|
|
if _, zero := rel.Field.ValueOf(obj); !zero {
|
|
rv := rel.Field.ReflectValueOf(obj)
|
|
if rv.Kind() != reflect.Ptr {
|
|
rv = rv.Addr()
|
|
}
|
|
|
|
for _, ref := range rel.References {
|
|
if ref.OwnPrimaryKey {
|
|
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
|
ref.ForeignKey.Set(rv, fv)
|
|
} else if ref.PrimaryValue != "" {
|
|
ref.ForeignKey.Set(rv, ref.PrimaryValue)
|
|
}
|
|
}
|
|
|
|
elems = reflect.Append(elems, rv)
|
|
}
|
|
}
|
|
|
|
if elems.Len() > 0 {
|
|
assignmentColumns := []string{}
|
|
for _, ref := range rel.References {
|
|
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
|
}
|
|
|
|
db.Session(&gorm.Session{}).Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: rel.FieldSchema.PrioritizedPrimaryField.DBName}},
|
|
DoUpdates: clause.AssignmentColumns(assignmentColumns),
|
|
}).Create(elems.Interface())
|
|
}
|
|
case reflect.Struct:
|
|
if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
|
|
f := rel.Field.ReflectValueOf(db.Statement.ReflectValue)
|
|
if f.Kind() != reflect.Ptr {
|
|
f = f.Addr()
|
|
}
|
|
|
|
assignmentColumns := []string{}
|
|
for _, ref := range rel.References {
|
|
if ref.OwnPrimaryKey {
|
|
fv, _ := ref.PrimaryKey.ValueOf(db.Statement.ReflectValue)
|
|
ref.ForeignKey.Set(f, fv)
|
|
} else if ref.PrimaryValue != "" {
|
|
ref.ForeignKey.Set(f, ref.PrimaryValue)
|
|
}
|
|
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
|
}
|
|
|
|
db.Session(&gorm.Session{}).Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: rel.FieldSchema.PrioritizedPrimaryField.DBName}},
|
|
DoUpdates: clause.AssignmentColumns(assignmentColumns),
|
|
}).Create(f.Interface())
|
|
}
|
|
}
|
|
}
|
|
|
|
// Save Has Many associations
|
|
for _, rel := range db.Statement.Schema.Relationships.HasMany {
|
|
if !saveAssociationCheck(db, rel, selectColumns, restricted) {
|
|
continue
|
|
}
|
|
|
|
fieldType := rel.Field.IndirectFieldType.Elem()
|
|
isPtr := fieldType.Kind() == reflect.Ptr
|
|
if !isPtr {
|
|
fieldType = reflect.PtrTo(fieldType)
|
|
}
|
|
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
|
appendToElems := func(v reflect.Value) {
|
|
if _, zero := rel.Field.ValueOf(v); !zero {
|
|
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
|
|
|
for i := 0; i < f.Len(); i++ {
|
|
elem := f.Index(i)
|
|
for _, ref := range rel.References {
|
|
if ref.OwnPrimaryKey {
|
|
pv, _ := ref.PrimaryKey.ValueOf(v)
|
|
ref.ForeignKey.Set(elem, pv)
|
|
} else if ref.PrimaryValue != "" {
|
|
ref.ForeignKey.Set(elem, ref.PrimaryValue)
|
|
}
|
|
}
|
|
|
|
if isPtr {
|
|
elems = reflect.Append(elems, elem)
|
|
} else {
|
|
elems = reflect.Append(elems, elem.Addr())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
switch db.Statement.ReflectValue.Kind() {
|
|
case reflect.Slice, reflect.Array:
|
|
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
|
appendToElems(db.Statement.ReflectValue.Index(i))
|
|
}
|
|
case reflect.Struct:
|
|
appendToElems(db.Statement.ReflectValue)
|
|
}
|
|
|
|
if elems.Len() > 0 {
|
|
assignmentColumns := []string{}
|
|
for _, ref := range rel.References {
|
|
assignmentColumns = append(assignmentColumns, ref.ForeignKey.DBName)
|
|
}
|
|
|
|
db.Session(&gorm.Session{}).Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: rel.FieldSchema.PrioritizedPrimaryField.DBName}},
|
|
DoUpdates: clause.AssignmentColumns(assignmentColumns),
|
|
}).Create(elems.Interface())
|
|
}
|
|
}
|
|
|
|
// Save Many2Many associations
|
|
for _, rel := range db.Statement.Schema.Relationships.Many2Many {
|
|
if !saveAssociationCheck(db, rel, selectColumns, restricted) {
|
|
continue
|
|
}
|
|
|
|
fieldType := rel.Field.IndirectFieldType.Elem()
|
|
isPtr := fieldType.Kind() == reflect.Ptr
|
|
if !isPtr {
|
|
fieldType = reflect.PtrTo(fieldType)
|
|
}
|
|
elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
|
|
joins := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(rel.JoinTable.ModelType)), 0, 0)
|
|
objs := []reflect.Value{}
|
|
|
|
appendToJoins := func(obj reflect.Value, elem reflect.Value) {
|
|
joinValue := reflect.New(rel.JoinTable.ModelType)
|
|
for _, ref := range rel.References {
|
|
if ref.OwnPrimaryKey {
|
|
fv, _ := ref.PrimaryKey.ValueOf(obj)
|
|
ref.ForeignKey.Set(joinValue, fv)
|
|
} else if ref.PrimaryValue != "" {
|
|
ref.ForeignKey.Set(joinValue, ref.PrimaryValue)
|
|
} else {
|
|
fv, _ := ref.PrimaryKey.ValueOf(elem)
|
|
ref.ForeignKey.Set(joinValue, fv)
|
|
}
|
|
}
|
|
joins = reflect.Append(joins, joinValue)
|
|
}
|
|
|
|
appendToElems := func(v reflect.Value) {
|
|
if _, zero := rel.Field.ValueOf(v); !zero {
|
|
f := reflect.Indirect(rel.Field.ReflectValueOf(v))
|
|
|
|
for i := 0; i < f.Len(); i++ {
|
|
elem := f.Index(i)
|
|
|
|
objs = append(objs, v)
|
|
if isPtr {
|
|
elems = reflect.Append(elems, elem)
|
|
} else {
|
|
elems = reflect.Append(elems, elem.Addr())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
switch db.Statement.ReflectValue.Kind() {
|
|
case reflect.Slice, reflect.Array:
|
|
for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
|
|
appendToElems(db.Statement.ReflectValue.Index(i))
|
|
}
|
|
case reflect.Struct:
|
|
appendToElems(db.Statement.ReflectValue)
|
|
}
|
|
|
|
if elems.Len() > 0 {
|
|
db.Session(&gorm.Session{}).Clauses(clause.OnConflict{DoNothing: true}).Create(elems.Interface())
|
|
|
|
for i := 0; i < elems.Len(); i++ {
|
|
appendToJoins(objs[i], elems.Index(i))
|
|
}
|
|
}
|
|
|
|
if joins.Len() > 0 {
|
|
db.Session(&gorm.Session{}).Clauses(clause.OnConflict{DoNothing: true}).Create(joins.Interface())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func saveAssociationCheck(db *gorm.DB, rel *schema.Relationship, selectColumns map[string]bool, restricted bool) bool {
|
|
savable := true
|
|
if value, ok := db.Get("gorm:save_association"); ok {
|
|
savable = utils.CheckTruth(value)
|
|
}
|
|
|
|
if savable {
|
|
if v, ok := selectColumns[rel.Name]; (ok && v) || (!ok && !restricted) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|