From 158bacefbef5d3995a81da02b238a5b3b8a3b024 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Sun, 19 Apr 2020 14:29:31 +0800
Subject: [PATCH] Add save has many relations

---
 callbacks/associations.go | 52 +++++++++++++++++++++++++++++++++++++
 tests/create.go           | 54 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/callbacks/associations.go b/callbacks/associations.go
index bbfbbc3d..6d976eac 100644
--- a/callbacks/associations.go
+++ b/callbacks/associations.go
@@ -99,6 +99,58 @@ func SaveAfterAssociations(db *gorm.DB) {
 			}
 		}
 	}
+
+	// Save Has Many associations
+	for _, rel := range db.Statement.Schema.Relationships.HasMany {
+		creatable, updatable, _ := saveAssociationCheck(db, rel.Field)
+		if !(creatable || updatable) {
+			continue
+		}
+
+		fieldType := rel.Field.IndirectFieldType.Elem()
+		isPtr := true
+		if fieldType.Kind() != reflect.Ptr {
+			isPtr = false
+			fieldType = reflect.PtrTo(fieldType)
+		}
+		elems := reflect.MakeSlice(reflect.SliceOf(fieldType), 0, 0)
+
+		switch db.Statement.ReflectValue.Kind() {
+		case reflect.Slice:
+			for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
+				db.Statement.ReflectValue.Index(i)
+			}
+		case reflect.Struct:
+			if _, zero := rel.Field.ValueOf(db.Statement.ReflectValue); !zero {
+				f := reflect.Indirect(rel.Field.ReflectValueOf(db.Statement.ReflectValue))
+
+				for i := 0; i < f.Len(); i++ {
+					elem := f.Index(i)
+					_, isZero := rel.FieldSchema.PrioritizedPrimaryField.ValueOf(elem)
+					for _, ref := range rel.References {
+						if ref.OwnPrimaryKey {
+							fv, _ := ref.PrimaryKey.ValueOf(db.Statement.ReflectValue)
+							ref.ForeignKey.Set(elem, fv)
+						} else if ref.PrimaryValue != "" {
+							ref.ForeignKey.Set(elem, ref.PrimaryValue)
+						}
+					}
+
+					if isZero && creatable {
+						if isPtr {
+							elems = reflect.Append(elems, elem)
+						} else {
+							elems = reflect.Append(elems, elem.Addr())
+						}
+					}
+				}
+			}
+		}
+
+		if elems.Len() > 0 {
+			db.Session(&gorm.Session{}).Create(elems.Interface())
+		}
+	}
 }
 
 func saveAssociationCheck(db *gorm.DB, field *schema.Field) (bool, bool, bool) {
diff --git a/tests/create.go b/tests/create.go
index 10b6b699..218e1e59 100644
--- a/tests/create.go
+++ b/tests/create.go
@@ -128,4 +128,58 @@ func TestCreateAssociations(t *testing.T, db *gorm.DB) {
 			}
 		}
 	})
+
+	t.Run("Create-HasManyAssociation", func(t *testing.T) {
+		var user = User{
+			Name:     "create",
+			Age:      18,
+			Birthday: Now(),
+			Pets:     []*Pet{{Name: "pet1"}, {Name: "pet2"}},
+		}
+
+		if err := db.Create(&user).Error; err != nil {
+			t.Fatalf("errors happened when create: %v", err)
+		}
+
+		for idx, pet := range user.Pets {
+			if pet.ID == 0 {
+				t.Fatalf("Failed to create pet #%v", idx)
+			}
+
+			var result Pet
+			db.First(&result, "id = ?", pet.ID)
+			if result.Name != pet.Name {
+				t.Errorf("Failed to query pet")
+			} else if result.UserID != user.ID {
+				t.Errorf("Failed to save relation")
+			}
+		}
+	})
+
+	t.Run("Create-HasManyAssociation-Polymorphic", func(t *testing.T) {
+		var user = User{
+			Name:     "create",
+			Age:      18,
+			Birthday: Now(),
+			Toys:     []Toy{{Name: "toy1"}, {Name: "toy2"}},
+		}
+
+		if err := db.Create(&user).Error; err != nil {
+			t.Fatalf("errors happened when create: %v", err)
+		}
+
+		for idx, toy := range user.Toys {
+			if toy.ID == 0 {
+				t.Fatalf("Failed to create toy #%v", idx)
+			}
+
+			var result Toy
+			db.First(&result, "id = ?", toy.ID)
+			if result.Name != toy.Name {
+				t.Errorf("Failed to query saved toy")
+			} else if result.OwnerID != fmt.Sprint(user.ID) || result.OwnerType != "users" {
+				t.Errorf("Failed to save relation")
+			}
+		}
+	})
 }