From 161be61b3d94bbecc0fd83704fa4c4d57a64a0a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=96=B9=E5=9C=A3=E5=8D=BF?= <fangshengqing@bytedance.com>
Date: Tue, 7 Nov 2023 21:53:45 +0800
Subject: [PATCH] create with schema map return pk

---
 callbacks/create.go | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/callbacks/create.go b/callbacks/create.go
index c10519c7..00bbea8b 100644
--- a/callbacks/create.go
+++ b/callbacks/create.go
@@ -153,7 +153,8 @@ func Create(config *Config) func(db *gorm.DB) {
 					db.AddError(db.Statement.Schema.PrioritizedPrimaryField.Set(db.Statement.Context, db.Statement.ReflectValue, insertID))
 				}
 			}
-		} else if db.Statement.Dest != nil {
+		}
+		if db.Statement.Dest != nil {
 			// append @id column with value for auto-increment primary key
 			// the @id value is correct, when: 1. without setting auto-increment primary key, 2. database AutoIncrementIncrement = 1
 			insertID, err := result.LastInsertId()
@@ -163,11 +164,16 @@ func Create(config *Config) func(db *gorm.DB) {
 				return
 			}
 
+			pkFieldName := "@id"
+			if db.Statement.Schema != nil && db.Statement.Schema.PrioritizedPrimaryField != nil {
+				pkFieldName = db.Statement.Schema.PrioritizedPrimaryField.DBName
+			}
+
 			switch values := db.Statement.Dest.(type) {
 			case map[string]interface{}:
-				values["@id"] = insertID
+				values[pkFieldName] = insertID
 			case *map[string]interface{}:
-				(*values)["@id"] = insertID
+				(*values)[pkFieldName] = insertID
 			case []map[string]interface{}, *[]map[string]interface{}:
 				mapValues, ok := values.([]map[string]interface{})
 				if !ok {
@@ -179,7 +185,7 @@ func Create(config *Config) func(db *gorm.DB) {
 				}
 				for _, mapValue := range mapValues {
 					if mapValue != nil {
-						mapValue["@id"] = insertID
+						mapValue[pkFieldName] = insertID
 					}
 					insertID += schema.DefaultAutoIncrementIncrement
 				}